From 3df70cdb25987d803ea362aab3693d6d00ce5a1d Mon Sep 17 00:00:00 2001
From: Daniel Ferrochio <daniel.ferrochio@telus.com>
Date: Thu, 16 May 2024 13:28:52 +0000
Subject: [PATCH] fix unittests

---
 .../tests/utils.py                                 | 14 ++++++++++++--
 1 file changed, 12 insertions(+), 2 deletions(-)

diff --git a/instrumentation/opentelemetry-instrumentation-confluent-kafka/tests/utils.py b/instrumentation/opentelemetry-instrumentation-confluent-kafka/tests/utils.py
index 0fc3bf696d..f16bb26b86 100644
--- a/instrumentation/opentelemetry-instrumentation-confluent-kafka/tests/utils.py
+++ b/instrumentation/opentelemetry-instrumentation-confluent-kafka/tests/utils.py
@@ -1,5 +1,5 @@
 from confluent_kafka import Consumer, Producer
-
+from typing import Optional
 
 class MockConsumer(Consumer):
     def __init__(self, queue, config):
@@ -20,11 +20,21 @@ def poll(self, timeout=None):
 
 
 class MockedMessage:
-    def __init__(self, topic: str, partition: int, offset: int, headers, key=None, value=None):
+    def __init__(
+            self,
+            topic: str,
+            partition: int,
+            offset: int,
+            headers,
+            key: Optional[str]=None,
+            value=Optional[str]=None
+        ):
         self._topic = topic
         self._partition = partition
         self._offset = offset
         self._headers = headers
+        self._key = key
+        self._value = value
 
     def topic(self):
         return self._topic