add test

Author: zhihali

instrumentation/opentelemetry-instrumentation-redis/tests/test_redis.py

@@ -12,14 +12,13 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import asyncio
-import logging
-from unittest import mock
-from unittest.mock import AsyncMock, patch
 
 import pytest
 import redis
-import redis.asyncio
+import redis.asyncio as redis_async
 from redis.exceptions import WatchError
+from unittest import mock, IsolatedAsyncioTestCase
+from unittest.mock import AsyncMock
 
 from opentelemetry import trace
 from opentelemetry.instrumentation.redis import RedisInstrumentor
@@ -29,7 +28,7 @@
     SpanAttributes,
 )
 from opentelemetry.test.test_base import TestBase
-from opentelemetry.trace import SpanKind, StatusCode
+from opentelemetry.trace import SpanKind
 
 
 class TestRedis(TestBase):
@@ -359,114 +358,32 @@ def test_successful_transaction(self):
         # For example, you might want to verify that the database index is correctly recorded
         self.assertEqual(span.attributes.get("db.redis.database_index"), 0)
 
-    def test_watch_error(self):
-        redis_client = redis.Redis()
-
-        # Mock the pipeline to raise a WatchError
-        mock_pipeline = mock.MagicMock()
-        mock_pipeline.watch.return_value = None
-        mock_pipeline.multi.return_value = mock_pipeline
-        mock_pipeline.execute.side_effect = WatchError("Watched variable changed")
-
-        with mock.patch.object(redis_client, "pipeline", return_value=mock_pipeline):
-            try:
-                with redis_client.pipeline() as pipe:
-                    pipe.watch("key")
-                    pipe.multi()
-                    pipe.set("key", "value")
-                    pipe.execute()
-            except WatchError:
-                pass  # We expect this exception to be raised
-
-        spans = self.memory_exporter.get_finished_spans()
-        self.assertEqual(len(spans), 1)
-        span = spans[0]
-
-        # Check that the span is not marked as an error
-        self.assertIsNone(span.status.status_code)
-
-        # Check that the WatchError is recorded as an event, not an exception
-        events = span.events
-        self.assertEqual(len(events), 1)
-        self.assertEqual(events[0].name, "exception")
-        self.assertEqual(events[0].attributes["exception.type"], "WatchError")
-        self.assertIn("Watched variable changed", events[0].attributes["exception.message"])
-
-        # Verify other span properties
-        self.assertEqual(span.name, "MULTI")
-        self.assertEqual(span.kind, SpanKind.CLIENT)
-        self.assertEqual(span.attributes.get("db.system"), "redis")
-
-
-import pytest
-import redis.asyncio
-from redis.exceptions import WatchError
-from opentelemetry import trace
-from opentelemetry.semconv.trace import SpanAttributes
-from opentelemetry.instrumentation.redis import RedisInstrumentor
-from opentelemetry.sdk.trace import TracerProvider
-from opentelemetry.sdk.trace.export import SimpleSpanProcessor
-from opentelemetry.sdk.trace.export.in_memory_span_exporter import \
-    InMemorySpanExporter  # This is the correct import for MemorySpanExporter
-
-
-class Test_Redis:
-    @pytest.fixture(autouse=True)
-    def setup_and_teardown(self):
-        # Setup
-        self.tracer_provider = TracerProvider()
-        self.memory_exporter = InMemorySpanExporter()
-        span_processor = SimpleSpanProcessor(self.memory_exporter)
-        self.tracer_provider.add_span_processor(span_processor)
-        trace.set_tracer_provider(self.tracer_provider)
 
+class TestRedisAsync(TestBase, IsolatedAsyncioTestCase):
+    def setUp(self):
+        super().setUp()
         RedisInstrumentor().instrument(tracer_provider=self.tracer_provider)
 
-        yield
-
-        # Teardown
+    def tearDown(self):
+        super().tearDown()
         RedisInstrumentor().uninstrument()
 
     @pytest.mark.asyncio
-    async def test_watch_error(self):
-        r = redis.asyncio.Redis()
-        await r.set("a", "0")
-
-        try:
-            async with r.pipeline(transaction=False) as pipe:
-                await pipe.watch("a")
-                a = await pipe.get("a")
-
-                # Simulate a change by another client
-                await r.set("a", "bad")
-
-                pipe.multi()
-                await pipe.set("a", str(int(a) + 1))
+    async def test_redis_operations(self):
+        async def redis_operations():
+            try:
+                r = redis_async.Redis()
+                async with r.pipeline(transaction=False) as pipe:
+                    await pipe.watch("a")
+                    await r.set("a", "bad")
+                    pipe.multi()
+                    await pipe.set("a", "1")
+                    await pipe.execute()
+            except WatchError:
+                pass
 
-                await pipe.execute()
-        except WatchError:
-            print("WatchError caught as expected")
-        else:
-            pytest.fail("WatchError was not raised")
+        await redis_operations()
 
         spans = self.memory_exporter.get_finished_spans()
-        assert len(spans) > 0, "No spans were recorded"
-
-        # Check the last span for WatchError evidence
-        last_span = spans[-1]
-
-        # The span itself should not be marked as an error
-        assert last_span.status.status_code is None
-
-        # Check for WatchError in span events
-        watch_error_events = [event for event in last_span.events
-                              if event.name == "exception" and
-                              event.attributes.get("exception.type") == "WatchError"]
-        assert len(watch_error_events) > 0, "WatchError event not found in span"
-
-        # Verify that the value in Redis wasn't changed due to the WatchError
-        final_value = await r.get("a")
-        assert final_value == b"bad"
-
-        # Clean up
-        await r.delete("a")
+        assert spans[-1].status.status_code == trace.StatusCode.UNSET
+        assert any(event.name == "WatchError" for event in spans[-1].events)