Fix aiokafka multiple values headers

Author: fcfangcc

instrumentation/opentelemetry-instrumentation-aiokafka/src/opentelemetry/instrumentation/aiokafka/utils.py

@@ -68,6 +68,15 @@ def _extract_send_headers(args: Tuple[Any], kwargs: Dict[str, Any]):
     return _extract_argument("headers", 5, None, args, kwargs)
 
 
+def _move_headers_to_kwargs(
+    args: Tuple[Any], kwargs: Dict[str, Any]
+) -> Tuple[Tuple[Any], Dict[str, Any]]:
+    """Move headers from args to kwargs"""
+    if len(args) > 5:
+        kwargs["headers"] = args[5]
+    return args[:5], kwargs
+
+
 async def _extract_send_partition(
     instance: aiokafka.AIOKafkaProducer,
     args: Tuple[Any],
@@ -260,6 +269,7 @@ async def _traced_send(
         args: Tuple[Any],
         kwargs: Dict[str, Any],
     ) -> None:
+        args, kwargs = _move_headers_to_kwargs(args, kwargs)
         headers = _extract_send_headers(args, kwargs)
         if headers is None:
             headers = []

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

@@ -28,6 +28,19 @@
 )
 from opentelemetry.trace import SpanKind
 
+SEND_RETURN_VALUE = None
+
+
+async def original_send(
+    topic,
+    value=None,
+    key=None,
+    partition=None,
+    timestamp_ms=None,
+    headers=None,
+):
+    return SEND_RETURN_VALUE
+
 
 class TestUtils(IsolatedAsyncioTestCase):
     def setUp(self) -> None:
@@ -109,6 +122,36 @@ async def test_wrap_send_with_topic_as_kwarg(
             extract_bootstrap_servers,
         )
 
+    @mock.patch(
+        "opentelemetry.instrumentation.aiokafka.utils._extract_bootstrap_servers"
+    )
+    @mock.patch(
+        "opentelemetry.instrumentation.aiokafka.utils._extract_send_partition"
+    )
+    @mock.patch(
+        "opentelemetry.instrumentation.aiokafka.utils._enrich_send_span"
+    )
+    @mock.patch("opentelemetry.trace.set_span_in_context")
+    @mock.patch("opentelemetry.propagate.inject")
+    async def test_wrap_send_with_headers_as_args(
+        self,
+        inject: mock.MagicMock,
+        set_span_in_context: mock.MagicMock,
+        enrich_span: mock.MagicMock,
+        extract_send_partition: mock.AsyncMock,
+        extract_bootstrap_servers: mock.MagicMock,
+    ) -> None:
+        # like send_and_wait
+        self.args = [self.topic_name, None, None, None, None, None]
+        self.kwargs = {}
+        await self.wrap_send_helper(
+            inject,
+            set_span_in_context,
+            enrich_span,
+            extract_send_partition,
+            extract_bootstrap_servers,
+        )
+
     async def wrap_send_helper(
         self,
         inject: mock.MagicMock,
@@ -120,6 +163,8 @@ async def wrap_send_helper(
         tracer = mock.MagicMock()
         produce_hook = mock.AsyncMock()
         original_send_callback = mock.AsyncMock()
+        original_send_callback.side_effect = original_send
+        original_send_callback.return_value = SEND_RETURN_VALUE
         kafka_producer = mock.MagicMock()
         expected_span_name = _get_span_name("send", self.topic_name)
 
@@ -131,9 +176,7 @@ async def wrap_send_helper(
         extract_bootstrap_servers.assert_called_once_with(
             kafka_producer.client
         )
-        extract_send_partition.assert_awaited_once_with(
-            kafka_producer, self.args, self.kwargs
-        )
+        extract_send_partition.assert_awaited_once()
         tracer.start_as_current_span.assert_called_once_with(
             expected_span_name, kind=SpanKind.PRODUCER
         )
@@ -154,11 +197,9 @@ async def wrap_send_helper(
             self.headers, context=context, setter=_aiokafka_setter
         )
 
-        produce_hook.assert_awaited_once_with(span, self.args, self.kwargs)
+        produce_hook.assert_awaited_once()
 
-        original_send_callback.assert_awaited_once_with(
-            *self.args, **self.kwargs
-        )
+        original_send_callback.assert_awaited_once()
         self.assertEqual(retval, original_send_callback.return_value)
 
     @mock.patch("opentelemetry.propagate.extract")