test(confluent-kafka): assert span and token cleanup on empty poll/consume

Author: sterchelen

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

@@ -314,6 +314,44 @@ def test_consume_empty_does_not_create_span(self) -> None:
         span_list = self.memory_exporter.get_finished_spans()
         self.assertEqual(len(span_list), 0)
 
+    def test_poll_empty_cleans_up_previous_span_and_token(self) -> None:
+        instrumentation = ConfluentKafkaInstrumentor()
+        consumer = MockConsumer(
+            [MockedMessage("topic-1", 0, 0, [])],
+            {
+                "bootstrap.servers": "localhost:29092",
+                "group.id": "mygroup",
+                "auto.offset.reset": "earliest",
+            },
+        )
+        consumer = instrumentation.instrument_consumer(consumer)
+        consumer.poll()  # non-empty: sets _current_consume_span and _current_context_token
+        self.assertIsNotNone(consumer._current_consume_span)
+        self.assertIsNotNone(consumer._current_context_token)
+
+        consumer.poll()  # empty: should clean up both
+        self.assertIsNone(consumer._current_consume_span)
+        self.assertIsNone(consumer._current_context_token)
+
+    def test_consume_empty_cleans_up_previous_span_and_token(self) -> None:
+        instrumentation = ConfluentKafkaInstrumentor()
+        consumer = MockConsumer(
+            [MockedMessage("topic-1", 0, 0, [])],
+            {
+                "bootstrap.servers": "localhost:29092",
+                "group.id": "mygroup",
+                "auto.offset.reset": "earliest",
+            },
+        )
+        consumer = instrumentation.instrument_consumer(consumer)
+        consumer.consume(1)  # non-empty: sets _current_consume_span and _current_context_token
+        self.assertIsNotNone(consumer._current_consume_span)
+        self.assertIsNotNone(consumer._current_context_token)
+
+        consumer.consume(1)  # empty: should clean up both
+        self.assertIsNone(consumer._current_consume_span)
+        self.assertIsNone(consumer._current_context_token)
+
     def test_close(self) -> None:
         instrumentation = ConfluentKafkaInstrumentor()
         mocked_messages = [