fix(kafka): wait_for_logs in kafka container to reduce lib requirement (testcontainers#377)

gudjonragnar · Gudjon Ragnar Brynjarsson · bearrito · commit 0cc9a72b1ad2 · 2024-03-30T18:00:48.000-04:00
Use `wait_for_logs` to wait for startup instead of waiting for successful connection via kafka-python. Also removes the dependency on kafka-python. Closes testcontainers#351 --------- Co-authored-by: Gudjon Ragnar Brynjarsson <gudjon.brynjarsson@controlant.com>
diff --git a/modules/kafka/testcontainers/kafka/__init__.py b/modules/kafka/testcontainers/kafka/__init__.py
@@ -3,11 +3,9 @@
 from io import BytesIO
 from textwrap import dedent
 
-from kafka import KafkaConsumer
-from kafka.errors import KafkaError, NoBrokersAvailable, UnrecognizedBrokerVersion
 from testcontainers.core.container import DockerContainer
 from testcontainers.core.utils import raise_for_deprecated_parameter
-from testcontainers.core.waiting_utils import wait_container_is_ready
+from testcontainers.core.waiting_utils import wait_for_logs
 
 
 class KafkaContainer(DockerContainer):
@@ -47,13 +45,6 @@ def get_bootstrap_server(self) -> str:
         port = self.get_exposed_port(self.port)
         return f"{host}:{port}"
 
-    @wait_container_is_ready(UnrecognizedBrokerVersion, NoBrokersAvailable, KafkaError, ValueError)
-    def _connect(self) -> None:
-        bootstrap_server = self.get_bootstrap_server()
-        consumer = KafkaConsumer(group_id="test", bootstrap_servers=[bootstrap_server])
-        if not consumer.bootstrap_connected():
-            raise KafkaError("Unable to connect with kafka container!")
-
     def tc_start(self) -> None:
         host = self.get_container_host_ip()
         port = self.get_exposed_port(self.port)
@@ -78,13 +69,13 @@ def tc_start(self) -> None:
         )
         self.create_file(data, KafkaContainer.TC_START_SCRIPT)
 
-    def start(self) -> "KafkaContainer":
+    def start(self, timeout=30) -> "KafkaContainer":
         script = KafkaContainer.TC_START_SCRIPT
         command = f'sh -c "while [ ! -f {script} ]; do sleep 0.1; done; sh {script}"'
         self.with_command(command)
         super().start()
         self.tc_start()
-        self._connect()
+        wait_for_logs(self, r".*\[KafkaServer id=\d+\] started.*", timeout=timeout)
         return self
 
     def create_file(self, content: bytes, path: str) -> None:
diff --git a/poetry.lock b/poetry.lock
diff --git a/pyproject.toml b/pyproject.toml
@@ -71,7 +71,6 @@ influxdb = { version = "*", optional = true }
 influxdb-client = { version = "*", optional = true }
 kubernetes = { version = "*", optional = true }
 pyyaml = { version = "*", optional = true }
-kafka-python = { version = "*", optional = true }
 python-keycloak = { version = "*", optional = true }
 boto3 = { version = "*", optional = true }
 minio = { version = "*", optional = true }
@@ -94,7 +93,7 @@ elasticsearch = []
 google = ["google-cloud-pubsub"]
 influxdb = ["influxdb", "influxdb-client"]
 k3s = ["kubernetes", "pyyaml"]
-kafka = ["kafka-python"]
+kafka = []
 keycloak = ["python-keycloak"]
 localstack = ["boto3"]
 minio = ["minio"]
@@ -123,7 +122,7 @@ anyio = "^4.3.0"
 psycopg2-binary = "*"
 pg8000 = "*"
 sqlalchemy = "*"
-
+kafka-python = "^2.0.2"
 
 [[tool.poetry.source]]
 name = "PyPI"
