separate assert function

bourbonkk · Oct 12, 2024 · 9ae469d · 9ae469d
1 parent b4af64c
commit 9ae469d
Showing 1 changed file with 10 additions and 6 deletions.
diff --git a/instrumentation/opentelemetry-instrumentation-confluent-kafka/tests/test_instrumentation.py b/instrumentation/opentelemetry-instrumentation-confluent-kafka/tests/test_instrumentation.py
@@ -284,15 +284,15 @@ def _compare_spans(self, spans, expected_spans):
                     expected_attribute_value, span.attributes[attribute_key]
                 )
 
-    def _assert_topic(self, expected_topic: str) -> None:
-        span_list = self.memory_exporter.get_finished_spans()
-        self.assertEqual(len(span_list), 1)
-        span = span_list[0]
+    def _assert_topic(self, span, expected_topic: str) -> None:
         self.assertEqual(
             span.attributes[SpanAttributes.MESSAGING_DESTINATION],
             expected_topic,
         )
 
+    def _assert_span_count(self, span_list, expected_count: int) -> None:
+        self.assertEqual(len(span_list), expected_count)
+
     def test_producer_poll(self) -> None:
         instrumentation = ConfluentKafkaInstrumentor()
         message_queue = []
@@ -308,7 +308,9 @@ def test_producer_poll(self) -> None:
         producer.produce(topic="topic-1", key="key-1", value="value-1")
         msg = producer.poll()
         self.assertIsNotNone(msg)
-        self._assert_topic("topic-1")
+        span_list = self.memory_exporter.get_finished_spans()
+        self._assert_span_count(span_list, 1)
+        self._assert_topic(span_list[0], "topic-1")
 
     def test_producer_flush(self) -> None:
         instrumentation = ConfluentKafkaInstrumentor()
@@ -325,4 +327,6 @@ def test_producer_flush(self) -> None:
         producer.produce(topic="topic-1", key="key-1", value="value-1")
         msg = producer.flush()
         self.assertIsNotNone(msg)
-        self._assert_topic("topic-1")
+        span_list = self.memory_exporter.get_finished_spans()
+        self._assert_span_count(span_list, 1)
+        self._assert_topic(span_list[0], "topic-1")