more tests

epsagon · Apr 8, 2024 · 578197b · 578197b
1 parent 59d3b1b
commit 578197b
Showing 1 changed file with 30 additions and 3 deletions.
diff --git a/instrumentation/opentelemetry-instrumentation-pika/tests/test_pika_instrumentation.py b/instrumentation/opentelemetry-instrumentation-pika/tests/test_pika_instrumentation.py
@@ -22,7 +22,10 @@
 from opentelemetry.instrumentation.pika.pika_instrumentor import (
     _consumer_callback_attribute_name,
 )
-from opentelemetry.instrumentation.pika.utils import dummy_callback
+from opentelemetry.instrumentation.pika.utils import (
+    ReadyMessagesDequeProxy,
+    dummy_callback,
+)
 from opentelemetry.trace import Tracer
 
 
@@ -42,15 +45,22 @@ def test_instrument_api(self) -> None:
             isinstance(BlockingConnection.channel, BoundFunctionWrapper)
         )
         self.assertTrue(
-            isinstance(_QueueConsumerGeneratorInfo.__init__, BoundFunctionWrapper)
+            isinstance(
+                _QueueConsumerGeneratorInfo.__init__, BoundFunctionWrapper
+            )
         )
         assert hasattr(
             instrumentation, "__opentelemetry_tracer_provider"
         ), "Tracer not stored for the object!"
-        instrumentation.uninstrument(channel=self.channel)
+        instrumentation.uninstrument()
         self.assertFalse(
             isinstance(BlockingConnection.channel, BoundFunctionWrapper)
         )
+        self.assertFalse(
+            isinstance(
+                _QueueConsumerGeneratorInfo.__init__, BoundFunctionWrapper
+            )
+        )
 
     @mock.patch(
         "opentelemetry.instrumentation.pika.PikaInstrumentor._instrument_channel_functions"
@@ -114,6 +124,23 @@ def test_instrument_basic_publish(
             self.channel.basic_publish, decorate_basic_publish.return_value
         )
 
+    def test_instrument_queue_consumer_generator(self) -> None:
+        instrumentation = PikaInstrumentor()
+        instrumentation.instrument()
+        generator_info = _QueueConsumerGeneratorInfo(
+            params=("queue", False, False), consumer_tag="tag"
+        )
+        self.assertTrue(
+            isinstance(generator_info.pending_events, ReadyMessagesDequeProxy)
+        )
+        instrumentation.uninstrument()
+        generator_info = _QueueConsumerGeneratorInfo(
+            params=("queue", False, False), consumer_tag="tag"
+        )
+        self.assertFalse(
+            isinstance(generator_info.pending_events, ReadyMessagesDequeProxy)
+        )
+
     def test_uninstrument_channel_functions(self) -> None:
         original_function = self.channel.basic_publish
         self.channel.basic_publish = mock.MagicMock()