Pass in config to client (#13970)

DataDog · Feb 15, 2023 · e2caa03 · e2caa03
1 parent 9585e2a
commit e2caa03
Show file tree

Hide file tree

Showing 3 changed files with 7 additions and 5 deletions.
diff --git a/kafka_consumer/datadog_checks/kafka_consumer/client/kafka_client_factory.py b/kafka_consumer/datadog_checks/kafka_consumer/client/kafka_client_factory.py
@@ -2,5 +2,5 @@
 from datadog_checks.kafka_consumer.client.kafka_python_client import KafkaPythonClient
 
 
-def make_client(check) -> KafkaClient:
-    return KafkaPythonClient(check)
+def make_client(check, config) -> KafkaClient:
+    return KafkaPythonClient(check, config)
diff --git a/kafka_consumer/datadog_checks/kafka_consumer/client/kafka_python_client.py b/kafka_consumer/datadog_checks/kafka_consumer/client/kafka_python_client.py
@@ -30,8 +30,9 @@ def token(self):
 
 
 class KafkaPythonClient(KafkaClient):
-    def __init__(self, check) -> None:
+    def __init__(self, check, config) -> None:
         self.check = check
+        self.config = config
         self.log = check.log
         self._kafka_client = None
         self._highwater_offsets = {}

diff --git a/kafka_consumer/datadog_checks/kafka_consumer/kafka_consumer.py b/kafka_consumer/datadog_checks/kafka_consumer/kafka_consumer.py
@@ -5,10 +5,11 @@
 from datadog_checks.base import AgentCheck, is_affirmative
 from datadog_checks.kafka_consumer.client.kafka_client_factory import make_client
 
+from .config_models import ConfigMixin
 from .constants import BROKER_REQUESTS_BATCH_SIZE, CONTEXT_UPPER_BOUND
 
 
-class KafkaCheck(AgentCheck):
+class KafkaCheck(AgentCheck, ConfigMixin):
 
     __NAMESPACE__ = 'kafka'
 
@@ -33,7 +34,7 @@ def __init__(self, name, init_config, instances):
         )
         self._consumer_groups = self.instance.get('consumer_groups', {})
         self._broker_requests_batch_size = self.instance.get('broker_requests_batch_size', BROKER_REQUESTS_BATCH_SIZE)
-        self.client = make_client(self)
+        self.client = make_client(self, self.config)
 
     def check(self, _):
         """The main entrypoint of the check."""