adding kafka test container support (#127)

* adding kafka containers support

* Removing the assertions inside the for loop as they are not guarenteed to provide correct feedback

Author: ash1425

requirements.in

@@ -1,4 +1,4 @@
--e file:.[docker-compose,mysql,oracle,postgresql,selenium,google-cloud-pubsub,mongo,redis,mssqlserver,neo4j]
+-e file:.[docker-compose,mysql,oracle,postgresql,selenium,google-cloud-pubsub,mongo,redis,mssqlserver,neo4j,kafka]
 codecov>=2.1.0
 flake8
 pytest

requirements/3.6.txt

@@ -91,6 +91,7 @@ jinja2==2.11.2
     # via sphinx
 jsonschema==3.2.0
     # via docker-compose
+kafka-python==2.0.2       # via testcontainers
 markupsafe==1.1.1
     # via jinja2
 mccabe==0.6.1

requirements/3.7.txt

@@ -91,6 +91,7 @@ jinja2==2.11.2
     # via sphinx
 jsonschema==3.2.0
     # via docker-compose
+kafka-python==2.0.2       # via testcontainers
 markupsafe==1.1.1
     # via jinja2
 mccabe==0.6.1

requirements/3.8.txt

@@ -85,6 +85,7 @@ jinja2==2.11.2
     # via sphinx
 jsonschema==3.2.0
     # via docker-compose
+kafka-python==2.0.2       # via testcontainers
 markupsafe==1.1.1
     # via jinja2
 mccabe==0.6.1

setup.py

@@ -62,7 +62,8 @@
         'mongo': ['pymongo'],
         'redis': ['redis'],
         'mssqlserver': ['pyodbc'],
-        'neo4j': ['neo4j']
+        'neo4j': ['neo4j'],
+        'kafka': ['kafka-python']
     },
     long_description_content_type="text/x-rst",
     long_description=long_description,

testcontainers/kafka.py

@@ -0,0 +1,84 @@
+import tarfile
+import time
+from io import BytesIO
+from textwrap import dedent
+
+from kafka import KafkaConsumer
+from kafka.errors import KafkaError
+
+from testcontainers.core.container import DockerContainer
+from testcontainers.core.waiting_utils import wait_container_is_ready
+
+
+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')
+
+    def get_bootstrap_server(self):
+        host = self.get_container_host_ip()
+        port = self.get_exposed_port(self.port_to_expose)
+        return '{}:{}'.format(host, port)
+
+    @wait_container_is_ready()
+    def _connect(self):
+        bootstrap_server = self.get_bootstrap_server()
+        consumer = KafkaConsumer(group_id='test', bootstrap_servers=[bootstrap_server])
+        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 = 'sh -c "while [ ! -f {} ]; do sleep 0.1; done; sh {}"'.format(script, script)
+        self.with_command(command)
+        super().start()
+        self.tc_start()
+        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)

tests/test_kafka.py

@@ -0,0 +1,31 @@
+from kafka import KafkaConsumer, KafkaProducer, TopicPartition
+
+from testcontainers.kafka import KafkaContainer
+
+
+def test_kafka_producer_consumer():
+    with KafkaContainer() as container:
+        produce_and_consume_kafka_message(container)
+
+
+def test_kafka_producer_consumer_custom_port():
+    with KafkaContainer(port_to_expose=9888) as container:
+        assert container.port_to_expose == 9888
+        produce_and_consume_kafka_message(container)
+
+
+def produce_and_consume_kafka_message(container):
+    topic = 'test-topic'
+    bootstrap_server = container.get_bootstrap_server()
+
+    producer = KafkaProducer(bootstrap_servers=[bootstrap_server])
+    producer.send(topic, b"verification message")
+    producer.flush()
+    producer.close()
+
+    consumer = KafkaConsumer(bootstrap_servers=[bootstrap_server])
+    tp = TopicPartition(topic, 0)
+    consumer.assign([tp])
+    consumer.seek_to_beginning()
+    assert consumer.end_offsets([tp])[tp] == 1, \
+        "Expected exactly one test message to be present on test topic !"