From 9ae469d8234dfc5a80ce5ab53c368776eeb0663f Mon Sep 17 00:00:00 2001 From: "allen.k1m" Date: Sat, 12 Oct 2024 14:03:11 +0900 Subject: [PATCH] separate assert function --- .../tests/test_instrumentation.py | 16 ++++++++++------ 1 file changed, 10 insertions(+), 6 deletions(-) diff --git a/instrumentation/opentelemetry-instrumentation-confluent-kafka/tests/test_instrumentation.py b/instrumentation/opentelemetry-instrumentation-confluent-kafka/tests/test_instrumentation.py index 7dad228001..986116900d 100644 --- 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")