Use wait_for_logs

Author: Gudjon Ragnar Brynjarsson

kafka/setup.py

@@ -12,7 +12,6 @@
     url="https://github.com/testcontainers/testcontainers-python",
     install_requires=[
         "testcontainers-core",
-        "kafka-python",
     ],
     python_requires=">=3.7",
 )

kafka/testcontainers/kafka/__init__.py

@@ -3,12 +3,9 @@
 from io import BytesIO
 from textwrap import dedent
 
-from kafka import KafkaConsumer
-from kafka.errors import KafkaError, UnrecognizedBrokerVersion, NoBrokersAvailable
-
 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):
@@ -49,13 +46,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)
@@ -80,13 +70,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: