From 7b24b8d65aaffec25dce8b8db7195ede9a486fbd Mon Sep 17 00:00:00 2001 From: Gudjon Ragnar Brynjarsson Date: Sun, 3 Mar 2024 18:16:01 +0000 Subject: [PATCH 1/3] Add redpanda testcontainer module --- modules/redpanda/README.rst | 1 + .../testcontainers/redpanda/__init__.py | 82 +++++++++++++++++++ modules/redpanda/tests/test_redpanda.py | 56 +++++++++++++ pyproject.toml | 3 + 4 files changed, 142 insertions(+) create mode 100644 modules/redpanda/README.rst create mode 100644 modules/redpanda/testcontainers/redpanda/__init__.py create mode 100644 modules/redpanda/tests/test_redpanda.py diff --git a/modules/redpanda/README.rst b/modules/redpanda/README.rst new file mode 100644 index 000000000..9e2962e80 --- /dev/null +++ b/modules/redpanda/README.rst @@ -0,0 +1 @@ +.. autoclass:: testcontainers.redpanda.RedpandaContainer diff --git a/modules/redpanda/testcontainers/redpanda/__init__.py b/modules/redpanda/testcontainers/redpanda/__init__.py new file mode 100644 index 000000000..a800026a1 --- /dev/null +++ b/modules/redpanda/testcontainers/redpanda/__init__.py @@ -0,0 +1,82 @@ +import tarfile +import time +from io import BytesIO +from textwrap import dedent + +from testcontainers.core.container import DockerContainer +from testcontainers.core.waiting_utils import wait_for_logs + + +class RedpandaContainer(DockerContainer): + """ + Redpanda container. + + Example: + + .. doctest:: + + >>> from testcontainers.redpanda import RedpandaContainer + + >>> with RedpandaContainer() as redpanda: + ... connection = redpanda.get_bootstrap_server() + """ + + TC_START_SCRIPT = "/tc-start.sh" + + def __init__( + self, + image: str = "docker.redpanda.com/redpandadata/redpanda:v23.1.13", + **kwargs, + ) -> None: + kwargs["entrypoint"] = "sh" + super(RedpandaContainer, self).__init__(image, **kwargs) + self.redpanda_port = 9092 + self.schema_registry_port = 8081 + self.with_exposed_ports(self.redpanda_port, self.schema_registry_port) + + def get_bootstrap_server(self) -> str: + host = self.get_container_host_ip() + port = self.get_exposed_port(self.redpanda_port) + return f"{host}:{port}" + + def get_schema_registry_address(self) -> str: + host = self.get_container_host_ip() + port = self.get_exposed_port(self.schema_registry_port) + return f"http://{host}:{port}" + + def tc_start(self) -> None: + host = self.get_container_host_ip() + port = self.get_exposed_port(self.redpanda_port) + + data = ( + dedent( + f""" + #!/bin/bash + /usr/bin/rpk redpanda start --mode dev-container --smp 1 --memory 1G \ + --kafka-addr PLAINTEXT://0.0.0.0:29092,OUTSIDE://0.0.0.0:9092 \ + --advertise-kafka-addr PLAINTEXT://127.0.0.1:29092,OUTSIDE://{host}:{port} + """ + ) + .strip() + .encode("utf-8") + ) + + self.create_file(data, RedpandaContainer.TC_START_SCRIPT) + + def start(self, timeout=10) -> "RedpandaContainer": + script = RedpandaContainer.TC_START_SCRIPT + command = f'-c "while [ ! -f {script} ]; do sleep 0.1; done; sh {script}"' + self.with_command(command) + super().start() + self.tc_start() + wait_for_logs(self, r".*Started Kafka API server.*", timeout=timeout) + return self + + def create_file(self, content: bytes, path: str) -> None: + 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) diff --git a/modules/redpanda/tests/test_redpanda.py b/modules/redpanda/tests/test_redpanda.py new file mode 100644 index 000000000..b65f82aa5 --- /dev/null +++ b/modules/redpanda/tests/test_redpanda.py @@ -0,0 +1,56 @@ +import requests +import json +from kafka import KafkaConsumer, KafkaProducer, TopicPartition, KafkaAdminClient +from kafka.admin import NewTopic +from testcontainers.redpanda import RedpandaContainer + + +def test_redpanda_producer_consumer(): + with RedpandaContainer() as container: + produce_and_consume_message(container) + + +def test_redpanda_confluent_latest(): + with RedpandaContainer( + image="docker.redpanda.com/redpandadata/redpanda:latest" + ) as container: + produce_and_consume_message(container) + + +def test_schema_registry(): + with RedpandaContainer() as container: + address = container.get_schema_registry_address() + subject_name = "test-subject-value" + url = f"{address}/subjects" + + payload = {"schema": json.dumps({"type": "string"})} + headers = {"Content-Type": "application/vnd.schemaregistry.v1+json"} + create_result = requests.post( + f"{url}/{subject_name}/versions", data=json.dumps(payload), headers=headers + ) + assert create_result.status_code == 200 + + result = requests.get(url) + assert result.status_code == 200 + assert subject_name in result.json() + + +def produce_and_consume_message(container): + topic = "test-topic" + bootstrap_server = container.get_bootstrap_server() + + admin = KafkaAdminClient(bootstrap_servers=[bootstrap_server]) + admin.create_topics([NewTopic(topic, 1, 1)]) + + producer = KafkaProducer(bootstrap_servers=[bootstrap_server]) + future = producer.send(topic, b"verification message") + future.get(timeout=10) + 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 !" diff --git a/pyproject.toml b/pyproject.toml index 7afb4cd96..ba6d5cd83 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -48,6 +48,7 @@ packages = [ { include = "testcontainers", from = "modules/postgres" }, { include = "testcontainers", from = "modules/rabbitmq" }, { include = "testcontainers", from = "modules/redis" }, + { include = "testcontainers", from = "modules/redpanda" }, { include = "testcontainers", from = "modules/selenium" } ] @@ -82,6 +83,7 @@ cx_Oracle = { version = "*", optional = true } psycopg2-binary = { version = "*", optional = true } pika = { version = "*", optional = true } redis = { version = "*", optional = true } +redpanda = { version = "*", optional = true } selenium = { version = "*", optional = true } [tool.poetry.extras] @@ -105,6 +107,7 @@ oracle = ["sqlalchemy", "cx_Oracle"] postgres = ["sqlalchemy", "psycopg2-binary"] rabbitmq = ["pika"] redis = ["redis"] +redpanda = [] selenium = ["selenium"] [tool.poetry.group.dev.dependencies] From 1e645a6060c9402d9940fd83a98986657da25f91 Mon Sep 17 00:00:00 2001 From: Dave Ankin Date: Fri, 8 Mar 2024 03:08:07 -0500 Subject: [PATCH 2/3] combine per feedback, not sure this is correct tho --- modules/kafka/README.rst | 1 + .../__init__.py => kafka/testcontainers/kafka/redpanda.py} | 0 modules/{redpanda => kafka}/tests/test_redpanda.py | 0 modules/redpanda/README.rst | 1 - 4 files changed, 1 insertion(+), 1 deletion(-) rename modules/{redpanda/testcontainers/redpanda/__init__.py => kafka/testcontainers/kafka/redpanda.py} (100%) rename modules/{redpanda => kafka}/tests/test_redpanda.py (100%) delete mode 100644 modules/redpanda/README.rst diff --git a/modules/kafka/README.rst b/modules/kafka/README.rst index 144c0fc2a..a481c7870 100644 --- a/modules/kafka/README.rst +++ b/modules/kafka/README.rst @@ -1,2 +1,3 @@ .. autoclass:: testcontainers.kafka.KafkaContainer .. title:: testcontainers.kafka.KafkaContainer +.. autoclass:: testcontainers.redpanda.RedpandaContainer diff --git a/modules/redpanda/testcontainers/redpanda/__init__.py b/modules/kafka/testcontainers/kafka/redpanda.py similarity index 100% rename from modules/redpanda/testcontainers/redpanda/__init__.py rename to modules/kafka/testcontainers/kafka/redpanda.py diff --git a/modules/redpanda/tests/test_redpanda.py b/modules/kafka/tests/test_redpanda.py similarity index 100% rename from modules/redpanda/tests/test_redpanda.py rename to modules/kafka/tests/test_redpanda.py diff --git a/modules/redpanda/README.rst b/modules/redpanda/README.rst deleted file mode 100644 index 9e2962e80..000000000 --- a/modules/redpanda/README.rst +++ /dev/null @@ -1 +0,0 @@ -.. autoclass:: testcontainers.redpanda.RedpandaContainer From 74339653a793e3aa86fb8449cb096da79ffa9eb4 Mon Sep 17 00:00:00 2001 From: Dave Ankin Date: Fri, 8 Mar 2024 03:13:48 -0500 Subject: [PATCH 3/3] tmp --- modules/kafka/testcontainers/kafka/__init__.py | 1 + modules/kafka/tests/test_redpanda.py | 2 +- pyproject.toml | 3 --- 3 files changed, 2 insertions(+), 4 deletions(-) diff --git a/modules/kafka/testcontainers/kafka/__init__.py b/modules/kafka/testcontainers/kafka/__init__.py index 399839433..2d6b5f552 100644 --- a/modules/kafka/testcontainers/kafka/__init__.py +++ b/modules/kafka/testcontainers/kafka/__init__.py @@ -8,6 +8,7 @@ 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.kafka.redpanda import RedpandaContainer class KafkaContainer(DockerContainer): diff --git a/modules/kafka/tests/test_redpanda.py b/modules/kafka/tests/test_redpanda.py index b65f82aa5..5f53b64b1 100644 --- a/modules/kafka/tests/test_redpanda.py +++ b/modules/kafka/tests/test_redpanda.py @@ -2,7 +2,7 @@ import json from kafka import KafkaConsumer, KafkaProducer, TopicPartition, KafkaAdminClient from kafka.admin import NewTopic -from testcontainers.redpanda import RedpandaContainer +from testcontainers.kafka import RedpandaContainer def test_redpanda_producer_consumer(): diff --git a/pyproject.toml b/pyproject.toml index ba6d5cd83..7afb4cd96 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -48,7 +48,6 @@ packages = [ { include = "testcontainers", from = "modules/postgres" }, { include = "testcontainers", from = "modules/rabbitmq" }, { include = "testcontainers", from = "modules/redis" }, - { include = "testcontainers", from = "modules/redpanda" }, { include = "testcontainers", from = "modules/selenium" } ] @@ -83,7 +82,6 @@ cx_Oracle = { version = "*", optional = true } psycopg2-binary = { version = "*", optional = true } pika = { version = "*", optional = true } redis = { version = "*", optional = true } -redpanda = { version = "*", optional = true } selenium = { version = "*", optional = true } [tool.poetry.extras] @@ -107,7 +105,6 @@ oracle = ["sqlalchemy", "cx_Oracle"] postgres = ["sqlalchemy", "psycopg2-binary"] rabbitmq = ["pika"] redis = ["redis"] -redpanda = [] selenium = ["selenium"] [tool.poetry.group.dev.dependencies]