confluent-kafka: test auto-instrumentation code path

Author: xrmx

instrumentation/opentelemetry-instrumentation-confluent-kafka/tests/test_instrumentation.py

@@ -14,9 +14,9 @@
 
 # pylint: disable=no-name-in-module
 
-from confluent_kafka import Consumer, Producer
-
 from opentelemetry.instrumentation.confluent_kafka import (
+    AutoInstrumentedConsumer,
+    AutoInstrumentedProducer,
     ConfluentKafkaInstrumentor,
     ProxiedConsumer,
     ProxiedProducer,
@@ -41,6 +41,8 @@
 
 class TestConfluentKafka(TestBase):
     def test_instrument_api(self) -> None:
+        from confluent_kafka import Consumer, Producer  # noqa: PLC0415
+
         instrumentation = ConfluentKafkaInstrumentor()
 
         producer = Producer({"bootstrap.servers": "localhost:29092"})
@@ -51,16 +53,22 @@ def test_instrument_api(self) -> None:
         producer = instrumentation.uninstrument_producer(producer)
         self.assertEqual(producer.__class__, Producer)
 
-        producer = Producer({"bootstrap.servers": "localhost:29092"})
-        producer = instrumentation.instrument_producer(producer)
+        consumer = Consumer(
+            {
+                "bootstrap.servers": "localhost:29092",
+                "group.id": "mygroup",
+                "auto.offset.reset": "earliest",
+            }
+        )
 
-        self.assertEqual(producer.__class__, ProxiedProducer)
+        consumer = instrumentation.instrument_consumer(consumer)
+        self.assertEqual(consumer.__class__, ProxiedConsumer)
 
-        producer = instrumentation.uninstrument_producer(producer)
-        self.assertEqual(producer.__class__, Producer)
+        consumer = instrumentation.uninstrument_consumer(consumer)
+        self.assertEqual(consumer.__class__, Consumer)
 
         consumer = Consumer(
-            {
+            **{
                 "bootstrap.servers": "localhost:29092",
                 "group.id": "mygroup",
                 "auto.offset.reset": "earliest",
@@ -73,7 +81,37 @@ def test_instrument_api(self) -> None:
         consumer = instrumentation.uninstrument_consumer(consumer)
         self.assertEqual(consumer.__class__, Consumer)
 
+    def test_instrument_api_with_instrument(self) -> None:
+        ConfluentKafkaInstrumentor().instrument()
+
+        from confluent_kafka import Consumer, Producer  # noqa: PLC0415
+
+        producer = Producer({"bootstrap.servers": "localhost:29092"})
+        self.assertEqual(producer.__class__, AutoInstrumentedProducer)
+
+        consumer = Consumer(
+            {
+                "bootstrap.servers": "localhost:29092",
+                "group.id": "mygroup",
+                "auto.offset.reset": "earliest",
+            }
+        )
+        self.assertEqual(consumer.__class__, AutoInstrumentedConsumer)
+
+        consumer = Consumer(
+            **{
+                "bootstrap.servers": "localhost:29092",
+                "group.id": "mygroup",
+                "auto.offset.reset": "earliest",
+            }
+        )
+        self.assertEqual(consumer.__class__, AutoInstrumentedConsumer)
+
+        ConfluentKafkaInstrumentor().uninstrument()
+
     def test_consumer_commit_method_exists(self) -> None:
+        from confluent_kafka import Consumer  # noqa: PLC0415
+
         instrumentation = ConfluentKafkaInstrumentor()
 
         consumer = Consumer(