Updated tests

Author: antonpirker

tests/integrations/redis/asyncio/test_redis_asyncio.py

@@ -1,6 +1,7 @@
 import pytest
 
 from sentry_sdk import capture_message, start_transaction
+from sentry_sdk.consts import SPANDATA
 from sentry_sdk.integrations.redis import RedisIntegration
 
 from fakeredis.aioredis import FakeRedis
@@ -67,7 +68,13 @@ async def test_async_redis_pipeline(
         "redis.commands": {
             "count": 3,
             "first_ten": expected_first_ten,
-        }
+        },
+        SPANDATA.DB_SYSTEM: "redis",
+        SPANDATA.DB_NAME: "0",
+        SPANDATA.SERVER_ADDRESS: connection.connection_pool.connection_kwargs.get(
+            "host"
+        ),
+        SPANDATA.SERVER_PORT: 6379,
     }
     assert span["tags"] == {
         "redis.transaction": is_transaction,

tests/integrations/redis/test_redis.py

@@ -12,6 +12,14 @@
     import mock  # python < 3.3
 
 
+MOCK_CONNECTION_POOL = mock.MagicMock()
+MOCK_CONNECTION_POOL.connection_kwargs = {
+    "host": "localhost",
+    "port": 63791,
+    "db": 1,
+}
+
+
 def test_basic(sentry_init, capture_events):
     sentry_init(integrations=[RedisIntegration()])
     events = capture_events()
@@ -67,17 +75,9 @@ def test_redis_pipeline(
     (span,) = event["spans"]
     assert span["op"] == "db.redis"
     assert span["description"] == "redis.pipeline.execute"
-    assert span["data"] == {
-        "redis.commands": {
-            "count": 3,
-            "first_ten": expected_first_ten,
-        },
-        SPANDATA.DB_SYSTEM: "redis",
-        SPANDATA.DB_NAME: "0",
-        SPANDATA.SERVER_ADDRESS: connection.connection_pool.connection_kwargs.get(
-            "host"
-        ),
-        SPANDATA.SERVER_PORT: 6379,
+    assert span["data"]["redis.commands"] == {
+        "count": 3,
+        "first_ten": expected_first_ten,
     }
     assert span["tags"] == {
         "redis.transaction": is_transaction,
@@ -247,3 +247,51 @@ def test_breadcrumbs(sentry_init, capture_events):
         },
         "timestamp": crumbs[1]["timestamp"],
     }
+
+
+def test_db_connection_attributes_client(sentry_init, capture_events):
+    sentry_init(
+        traces_sample_rate=1.0,
+        integrations=[RedisIntegration()],
+    )
+    events = capture_events()
+
+    with start_transaction():
+        connection = FakeStrictRedis(connection_pool=MOCK_CONNECTION_POOL)
+        connection.get("foobar")
+
+    (event,) = events
+    (span,) = event["spans"]
+
+    assert span["op"] == "db.redis"
+    assert span["description"] == "GET 'foobar'"
+    assert span["data"][SPANDATA.DB_SYSTEM] == "redis"
+    assert span["data"][SPANDATA.DB_NAME] == "1"
+    assert span["data"][SPANDATA.SERVER_ADDRESS] == "localhost"
+    assert span["data"][SPANDATA.SERVER_PORT] == 63791
+
+
+def test_db_connection_attributes_pipeline(sentry_init, capture_events):
+    sentry_init(
+        traces_sample_rate=1.0,
+        integrations=[RedisIntegration()],
+    )
+    events = capture_events()
+
+    with start_transaction():
+        connection = FakeStrictRedis(connection_pool=MOCK_CONNECTION_POOL)
+        pipeline = connection.pipeline(transaction=False)
+        pipeline.get("foo")
+        pipeline.set("bar", 1)
+        pipeline.set("baz", 2)
+        pipeline.execute()
+
+    (event,) = events
+    (span,) = event["spans"]
+
+    assert span["op"] == "db.redis"
+    assert span["description"] == "redis.pipeline.execute"
+    assert span["data"][SPANDATA.DB_SYSTEM] == "redis"
+    assert span["data"][SPANDATA.DB_NAME] == "1"
+    assert span["data"][SPANDATA.SERVER_ADDRESS] == "localhost"
+    assert span["data"][SPANDATA.SERVER_PORT] == 63791

tests/integrations/rediscluster/test_rediscluster.py

@@ -1,11 +1,25 @@
 import pytest
+
 from sentry_sdk import capture_message
-from sentry_sdk.consts import SPANDATA
 from sentry_sdk.api import start_transaction
+from sentry_sdk.consts import SPANDATA
 from sentry_sdk.integrations.redis import RedisIntegration
 
+try:
+    from unittest import mock
+except ImportError:
+    import mock
+
 import rediscluster
-from fakeredis import FakeStrictRedis
+
+
+MOCK_CONNECTION_POOL = mock.MagicMock()
+MOCK_CONNECTION_POOL.connection_kwargs = {
+    "host": "localhost",
+    "port": 63791,
+    "db": 1,
+}
+
 
 rediscluster_classes = [rediscluster.RedisCluster]
 
@@ -20,7 +34,7 @@ def monkeypatch_rediscluster_classes(reset_integrations):
     except AttributeError:
         pipeline_cls = rediscluster.StrictClusterPipeline
     rediscluster.RedisCluster.pipeline = lambda *_, **__: pipeline_cls(
-        connection_pool=True
+        connection_pool=MOCK_CONNECTION_POOL
     )
     pipeline_cls.execute = lambda *_, **__: None
     for cls in rediscluster_classes:
@@ -32,7 +46,7 @@ def test_rediscluster_basic(rediscluster_cls, sentry_init, capture_events):
     sentry_init(integrations=[RedisIntegration()])
     events = capture_events()
 
-    rc = rediscluster_cls(connection_pool=True)
+    rc = rediscluster_cls(connection_pool=MOCK_CONNECTION_POOL)
     rc.get("foobar")
     capture_message("hi")
 
@@ -70,17 +84,7 @@ def test_rediscluster_pipeline(
     )
     events = capture_events()
 
-    connection = FakeStrictRedis()
-    startup_nodes = [
-        {
-            "host": connection.connection_pool.connection_kwargs.get("host"),
-            "port": connection.connection_pool.connection_kwargs.get("port"),
-        }
-    ]
-    print("////////////")
-    print(startup_nodes)
-
-    rc = rediscluster.RedisCluster(startup_nodes=startup_nodes)
+    rc = rediscluster.RedisCluster(connection_pool=MOCK_CONNECTION_POOL)
     with start_transaction():
         pipeline = rc.pipeline()
         pipeline.get("foo")
@@ -98,11 +102,66 @@ def test_rediscluster_pipeline(
             "first_ten": expected_first_ten,
         },
         SPANDATA.DB_SYSTEM: "redis",
-        SPANDATA.DB_NAME: "0",
-        SPANDATA.SERVER_ADDRESS: rc.connection_pool.connection_kwargs.get("host"),
-        SPANDATA.SERVER_PORT: 6379,
+        SPANDATA.DB_NAME: "1",
+        SPANDATA.SERVER_ADDRESS: "localhost",
+        SPANDATA.SERVER_PORT: 63791,
     }
     assert span["tags"] == {
         "redis.transaction": False,  # For Cluster, this is always False
         "redis.is_cluster": True,
     }
+
+
+@pytest.mark.parametrize("rediscluster_cls", rediscluster_classes)
+def test_db_connection_attributes_client(sentry_init, capture_events, rediscluster_cls):
+    sentry_init(
+        traces_sample_rate=1.0,
+        integrations=[RedisIntegration()],
+    )
+    events = capture_events()
+
+    rc = rediscluster_cls(connection_pool=MOCK_CONNECTION_POOL)
+    with start_transaction():
+        rc.get("foobar")
+
+    (event,) = events
+    (span,) = event["spans"]
+
+    assert span["data"] == {
+        SPANDATA.DB_SYSTEM: "redis",
+        SPANDATA.DB_NAME: "1",
+        SPANDATA.SERVER_ADDRESS: "localhost",
+        SPANDATA.SERVER_PORT: 63791,
+    }
+
+
+@pytest.mark.parametrize("rediscluster_cls", rediscluster_classes)
+def test_db_connection_attributes_pipeline(
+    sentry_init, capture_events, rediscluster_cls
+):
+    sentry_init(
+        traces_sample_rate=1.0,
+        integrations=[RedisIntegration()],
+    )
+    events = capture_events()
+
+    rc = rediscluster.RedisCluster(connection_pool=MOCK_CONNECTION_POOL)
+    with start_transaction():
+        pipeline = rc.pipeline()
+        pipeline.get("foo")
+        pipeline.execute()
+
+    (event,) = events
+    (span,) = event["spans"]
+    assert span["op"] == "db.redis"
+    assert span["description"] == "redis.pipeline.execute"
+    assert span["data"] == {
+        "redis.commands": {
+            "count": 1,
+            "first_ten": ["GET 'foo'"],
+        },
+        SPANDATA.DB_SYSTEM: "redis",
+        SPANDATA.DB_NAME: "1",
+        SPANDATA.SERVER_ADDRESS: "localhost",
+        SPANDATA.SERVER_PORT: 63791,
+    }