Simplify Kafka configuration.

Author: tillahoffmann

testcontainers/kafka.py

@@ -1,8 +1,3 @@
-import tarfile
-import time
-from io import BytesIO
-from textwrap import dedent
-
 from kafka import KafkaConsumer
 from kafka.errors import KafkaError, UnrecognizedBrokerVersion, NoBrokersAvailable
 
@@ -12,23 +7,28 @@
 
 class KafkaContainer(DockerContainer):
     KAFKA_PORT = 9093
-    TC_START_SCRIPT = '/tc-start.sh'
 
     def __init__(self, image="confluentinc/cp-kafka:5.4.3", port_to_expose=KAFKA_PORT):
         super(KafkaContainer, self).__init__(image)
         self.port_to_expose = port_to_expose
         self.with_exposed_ports(self.port_to_expose)
-        listeners = 'PLAINTEXT://0.0.0.0:{},BROKER://0.0.0.0:9092'.format(port_to_expose)
-        self.with_env('KAFKA_LISTENERS', listeners)
-        self.with_env('KAFKA_LISTENER_SECURITY_PROTOCOL_MAP',
-                      'BROKER:PLAINTEXT,PLAINTEXT:PLAINTEXT')
-        self.with_env('KAFKA_INTER_BROKER_LISTENER_NAME', 'BROKER')
 
-        self.with_env('KAFKA_BROKER_ID', '1')
-        self.with_env('KAFKA_OFFSETS_TOPIC_REPLICATION_FACTOR', '1')
-        self.with_env('KAFKA_OFFSETS_TOPIC_NUM_PARTITIONS', '1')
-        self.with_env('KAFKA_LOG_FLUSH_INTERVAL_MESSAGES', '10000000')
-        self.with_env('KAFKA_GROUP_INITIAL_REBALANCE_DELAY_MS', '0')
+        env = {
+            'KAFKA_LISTENERS': f'PLAINTEXT://0.0.0.0:{port_to_expose},BROKER://0.0.0.0:9092',
+            'KAFKA_LISTENER_SECURITY_PROTOCOL_MAP': 'BROKER:PLAINTEXT,PLAINTEXT:PLAINTEXT',
+            'KAFKA_INTER_BROKER_LISTENER_NAME': 'BROKER',
+            'KAFKA_BROKER_ID': '1',
+            'KAFKA_OFFSETS_TOPIC_REPLICATION_FACTOR': '1',
+            'KAFKA_OFFSETS_TOPIC_NUM_PARTITIONS': '1',
+            'KAFKA_LOG_FLUSH_INTERVAL_MESSAGES': '10000000',
+            'KAFKA_GROUP_INITIAL_REBALANCE_DELAY_MS': '0',
+            'KAFKA_ZOOKEEPER_CONNECT': 'localhost:2181',
+        }
+        for key, value in env.items():
+            self.with_env(key, value)
+
+        # Start zookeeper first because it doesn't need any port mapping information.
+        self.with_command("zookeeper-server-start /etc/kafka/zookeeper.properties")
 
     def get_bootstrap_server(self):
         host = self.get_container_host_ip()
@@ -42,43 +42,20 @@ def _connect(self):
         if not consumer.topics():
             raise KafkaError("Unable to connect with kafka container!")
 
-    def tc_start(self):
-        port = self.get_exposed_port(self.port_to_expose)
-        listeners = 'PLAINTEXT://localhost:{},BROKER://$(hostname -i):9092'.format(port)
-        data = (
-            dedent(
-                """
-                #!/bin/bash
-                echo 'clientPort=2181' > zookeeper.properties
-                echo 'dataDir=/var/lib/zookeeper/data' >> zookeeper.properties
-                echo 'dataLogDir=/var/lib/zookeeper/log' >> zookeeper.properties
-                zookeeper-server-start zookeeper.properties &
-                export KAFKA_ZOOKEEPER_CONNECT='localhost:2181'
-                export KAFKA_ADVERTISED_LISTENERS={}
-                . /etc/confluent/docker/bash-config
-                /etc/confluent/docker/configure
-                /etc/confluent/docker/launch
-                """.format(listeners)
-            )
-            .strip()
-            .encode('utf-8')
-        )
-        self.create_file(data, KafkaContainer.TC_START_SCRIPT)
-
     def start(self):
-        script = KafkaContainer.TC_START_SCRIPT
-        command = f'bash -c "while [ ! -f {script} ]; do sleep 0.1; done; bash {script}"'
-        self.with_command(command)
         super().start()
-        self.tc_start()
+        # Set the environment variables for which we need to know the port mappings and configure
+        # kafka.
+        exposed_port = self.get_exposed_port(self.port_to_expose)
+        host = self.get_container_host_ip()
+        advertised_listeners = f'PLAINTEXT://localhost:{exposed_port},BROKER://{host}:9092'
+        code, output = self._container.exec_run(
+            f'sh -c \'KAFKA_ADVERTISED_LISTENERS="{advertised_listeners}" '
+            '/etc/confluent/docker/configure\''
+        )
+        assert code == 0, output
+
+        # Start Kafka.
+        self._container.exec_run('/etc/confluent/docker/launch', detach=True)
         self._connect()
         return self
-
-    def create_file(self, content: bytes, path: str):
-        with BytesIO() as archive, tarfile.TarFile(fileobj=archive, mode="w") as tar:
-            tarinfo = tarfile.TarInfo(name=path)
-            tarinfo.size = len(content)
-            tarinfo.mtime = time.time()
-            tar.addfile(tarinfo, BytesIO(content))
-            archive.seek(0)
-            self.get_wrapped_container().put_archive("/", archive)