test: add test_wrap_getmany

Author: dimastbk

instrumentation/opentelemetry-instrumentation-aiokafka/tests/test_utils.py

@@ -15,6 +15,8 @@
 
 from unittest import IsolatedAsyncioTestCase, mock
 
+import aiokafka
+
 from opentelemetry.instrumentation.aiokafka.utils import (
     AIOKafkaContextGetter,
     AIOKafkaContextSetter,
@@ -23,6 +25,7 @@
     _create_consumer_span,
     _extract_send_partition,
     _get_span_name,
+    _wrap_getmany,
     _wrap_getone,
     _wrap_send,
 )
@@ -174,7 +177,7 @@ async def wrap_send_helper(
     @mock.patch(
         "opentelemetry.instrumentation.aiokafka.utils._extract_consumer_group"
     )
-    async def test_wrap_next(
+    async def test_wrap_getone(
         self,
         extract_consumer_group: mock.MagicMock,
         extract_client_id: mock.MagicMock,
@@ -184,12 +187,12 @@ async def test_wrap_next(
     ) -> None:
         tracer = mock.MagicMock()
         consume_hook = mock.AsyncMock()
-        original_next_callback = mock.AsyncMock()
+        original_getone_callback = mock.AsyncMock()
         kafka_consumer = mock.MagicMock()
 
-        wrapped_next = _wrap_getone(tracer, consume_hook)
-        record = await wrapped_next(
-            original_next_callback, kafka_consumer, self.args, self.kwargs
+        wrapped_getone = _wrap_getone(tracer, consume_hook)
+        record = await wrapped_getone(
+            original_getone_callback, kafka_consumer, self.args, self.kwargs
         )
 
         extract_bootstrap_servers.assert_called_once_with(
@@ -203,10 +206,10 @@ async def test_wrap_next(
         extract_consumer_group.assert_called_once_with(kafka_consumer)
         consumer_group = extract_consumer_group.return_value
 
-        original_next_callback.assert_awaited_once_with(
+        original_getone_callback.assert_awaited_once_with(
             *self.args, **self.kwargs
         )
-        self.assertEqual(record, original_next_callback.return_value)
+        self.assertEqual(record, original_getone_callback.return_value)
 
         extract.assert_called_once_with(
             record.headers, getter=_aiokafka_getter
@@ -225,6 +228,85 @@ async def test_wrap_next(
             self.kwargs,
         )
 
+    @mock.patch("opentelemetry.propagate.extract")
+    @mock.patch(
+        "opentelemetry.instrumentation.aiokafka.utils._create_consumer_span"
+    )
+    @mock.patch(
+        "opentelemetry.instrumentation.aiokafka.utils._enrich_getmany_topic_span"
+    )
+    @mock.patch(
+        "opentelemetry.instrumentation.aiokafka.utils._enrich_getmany_poll_span"
+    )
+    @mock.patch(
+        "opentelemetry.instrumentation.aiokafka.utils._extract_bootstrap_servers"
+    )
+    @mock.patch(
+        "opentelemetry.instrumentation.aiokafka.utils._extract_client_id"
+    )
+    @mock.patch(
+        "opentelemetry.instrumentation.aiokafka.utils._extract_consumer_group"
+    )
+    async def test_wrap_getmany(
+        self,
+        extract_consumer_group: mock.MagicMock,
+        extract_client_id: mock.MagicMock,
+        extract_bootstrap_servers: mock.MagicMock,
+        enrich_getmany_poll_span: mock.MagicMock,
+        enrich_getmany_topic_span: mock.MagicMock,
+        _create_consumer_span: mock.MagicMock,
+        extract: mock.MagicMock,
+    ) -> None:
+        tracer = mock.MagicMock()
+        consume_hook = mock.AsyncMock()
+        record_mock = mock.MagicMock()
+        original_getmany_callback = mock.AsyncMock(
+            return_value={
+                aiokafka.TopicPartition(topic="topic_1", partition=0): [
+                    record_mock
+                ]
+            }
+        )
+        kafka_consumer = mock.MagicMock()
+
+        wrapped_getmany = _wrap_getmany(tracer, consume_hook)
+        records = await wrapped_getmany(
+            original_getmany_callback, kafka_consumer, self.args, self.kwargs
+        )
+
+        extract_bootstrap_servers.assert_called_once_with(
+            kafka_consumer._client
+        )
+        bootstrap_servers = extract_bootstrap_servers.return_value
+
+        extract_client_id.assert_called_once_with(kafka_consumer._client)
+        client_id = extract_client_id.return_value
+
+        extract_consumer_group.assert_called_once_with(kafka_consumer)
+        consumer_group = extract_consumer_group.return_value
+
+        original_getmany_callback.assert_awaited_once_with(
+            *self.args, **self.kwargs
+        )
+        self.assertEqual(records, original_getmany_callback.return_value)
+
+        extract.assert_called_once_with(
+            record_mock.headers, getter=_aiokafka_getter
+        )
+        context = extract.return_value
+
+        _create_consumer_span.assert_called_once_with(
+            tracer,
+            consume_hook,
+            record_mock,
+            context,
+            bootstrap_servers,
+            client_id,
+            consumer_group,
+            self.args,
+            self.kwargs,
+        )
+
     @mock.patch("opentelemetry.trace.set_span_in_context")
     @mock.patch("opentelemetry.context.attach")
     @mock.patch(