|
| 1 | +# Licensed under the Apache License, Version 2.0 (the "License"); |
| 2 | +# you may not use this file except in compliance with the License. |
| 3 | +# You may obtain a copy of the License at |
| 4 | +# |
| 5 | +# http://www.apache.org/licenses/LICENSE-2.0 |
| 6 | +# |
| 7 | +# Unless required by applicable law or agreed to in writing, software |
| 8 | +# distributed under the License is distributed on an "AS IS" BASIS, |
| 9 | +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| 10 | +# See the License for the specific language governing permissions and |
| 11 | +# limitations under the License. |
| 12 | + |
| 13 | +""" |
| 14 | +Docker-managed Kafka broker for integration tests. |
| 15 | +
|
| 16 | +Provides automated image pull, idempotent container start, health check, |
| 17 | +topic lifecycle management, and data cleanup via KRaft-mode single-node Kafka. |
| 18 | +
|
| 19 | +Usage:: |
| 20 | +
|
| 21 | + mgr = KafkaDockerManager() |
| 22 | + mgr.setup_kafka() |
| 23 | + mgr.create_topics_if_not_exist(["input-topic", "output-topic"]) |
| 24 | + ... |
| 25 | + mgr.clear_all_topics() |
| 26 | + mgr.teardown_kafka() |
| 27 | +""" |
| 28 | + |
| 29 | +import logging |
| 30 | +import time |
| 31 | +from typing import List |
| 32 | + |
| 33 | +import docker |
| 34 | +from docker.errors import APIError, NotFound |
| 35 | +from confluent_kafka.admin import AdminClient, NewTopic |
| 36 | + |
| 37 | +logger = logging.getLogger(__name__) |
| 38 | + |
| 39 | +_DEFAULT_IMAGE = "apache/kafka:3.7.0" |
| 40 | +_DEFAULT_CONTAINER = "fs-integration-kafka-broker" |
| 41 | +_DEFAULT_BOOTSTRAP = "127.0.0.1:9092" |
| 42 | + |
| 43 | + |
| 44 | +class KafkaDockerManager: |
| 45 | + """ |
| 46 | + Manages a single-node Kafka broker inside a Docker container (KRaft mode). |
| 47 | +
|
| 48 | + The class is intentionally stateless with respect to topics: every public |
| 49 | + method is idempotent so that tests can call ``setup_kafka()`` multiple |
| 50 | + times without side-effects. |
| 51 | + """ |
| 52 | + |
| 53 | + def __init__( |
| 54 | + self, |
| 55 | + image: str = _DEFAULT_IMAGE, |
| 56 | + container_name: str = _DEFAULT_CONTAINER, |
| 57 | + bootstrap_servers: str = _DEFAULT_BOOTSTRAP, |
| 58 | + ) -> None: |
| 59 | + self.docker_client = docker.from_env() |
| 60 | + self.image_name = image |
| 61 | + self.container_name = container_name |
| 62 | + self.bootstrap_servers = bootstrap_servers |
| 63 | + |
| 64 | + # ------------------------------------------------------------------ |
| 65 | + # Full setup / teardown |
| 66 | + # ------------------------------------------------------------------ |
| 67 | + |
| 68 | + def setup_kafka(self) -> None: |
| 69 | + """Pull image -> start container -> wait for readiness.""" |
| 70 | + self._ensure_image() |
| 71 | + self._ensure_container() |
| 72 | + self._wait_for_readiness() |
| 73 | + |
| 74 | + def teardown_kafka(self) -> None: |
| 75 | + """Stop and remove the Kafka container.""" |
| 76 | + try: |
| 77 | + container = self.docker_client.containers.get(self.container_name) |
| 78 | + logger.info("Stopping Kafka container '%s' ...", self.container_name) |
| 79 | + container.stop() |
| 80 | + except NotFound: |
| 81 | + pass |
| 82 | + except APIError as exc: |
| 83 | + logger.warning("Error while stopping Kafka: %s", exc) |
| 84 | + |
| 85 | + # ------------------------------------------------------------------ |
| 86 | + # Image management |
| 87 | + # ------------------------------------------------------------------ |
| 88 | + |
| 89 | + def _ensure_image(self) -> None: |
| 90 | + try: |
| 91 | + self.docker_client.images.get(self.image_name) |
| 92 | + logger.info("Image '%s' already present locally.", self.image_name) |
| 93 | + except NotFound: |
| 94 | + logger.info("Pulling Kafka image '%s' ...", self.image_name) |
| 95 | + self.docker_client.images.pull(self.image_name) |
| 96 | + logger.info("Image pulled successfully.") |
| 97 | + |
| 98 | + # ------------------------------------------------------------------ |
| 99 | + # Container management (KRaft single-node, apache/kafka official image) |
| 100 | + # ------------------------------------------------------------------ |
| 101 | + |
| 102 | + def _ensure_container(self) -> None: |
| 103 | + try: |
| 104 | + container = self.docker_client.containers.get(self.container_name) |
| 105 | + if container.status != "running": |
| 106 | + logger.info( |
| 107 | + "Container '%s' exists but is not running; starting ...", |
| 108 | + self.container_name, |
| 109 | + ) |
| 110 | + container.start() |
| 111 | + else: |
| 112 | + logger.info( |
| 113 | + "Container '%s' is already running.", self.container_name |
| 114 | + ) |
| 115 | + except NotFound: |
| 116 | + logger.info( |
| 117 | + "Creating Kafka container '%s' ...", self.container_name |
| 118 | + ) |
| 119 | + env = { |
| 120 | + "KAFKA_NODE_ID": "1", |
| 121 | + "KAFKA_PROCESS_ROLES": "broker,controller", |
| 122 | + "KAFKA_CONTROLLER_LISTENER_NAMES": "CONTROLLER", |
| 123 | + "KAFKA_LISTENERS": ( |
| 124 | + "PLAINTEXT://0.0.0.0:9092,CONTROLLER://0.0.0.0:9093" |
| 125 | + ), |
| 126 | + "KAFKA_LISTENER_SECURITY_PROTOCOL_MAP": ( |
| 127 | + "CONTROLLER:PLAINTEXT,PLAINTEXT:PLAINTEXT" |
| 128 | + ), |
| 129 | + "KAFKA_ADVERTISED_LISTENERS": ( |
| 130 | + f"PLAINTEXT://{self.bootstrap_servers}" |
| 131 | + ), |
| 132 | + "KAFKA_CONTROLLER_QUORUM_VOTERS": "1@localhost:9093", |
| 133 | + "KAFKA_OFFSETS_TOPIC_REPLICATION_FACTOR": "1", |
| 134 | + "KAFKA_TRANSACTION_STATE_LOG_REPLICATION_FACTOR": "1", |
| 135 | + "KAFKA_TRANSACTION_STATE_LOG_MIN_ISR": "1", |
| 136 | + "CLUSTER_ID": "fs-integration-test-cluster-01", |
| 137 | + } |
| 138 | + self.docker_client.containers.run( |
| 139 | + image=self.image_name, |
| 140 | + name=self.container_name, |
| 141 | + ports={"9092/tcp": 9092}, |
| 142 | + environment=env, |
| 143 | + detach=True, |
| 144 | + remove=True, |
| 145 | + ) |
| 146 | + |
| 147 | + def _wait_for_readiness(self, timeout: int = 60) -> None: |
| 148 | + logger.info( |
| 149 | + "Waiting for Kafka to become ready at %s ...", |
| 150 | + self.bootstrap_servers, |
| 151 | + ) |
| 152 | + deadline = time.time() + timeout |
| 153 | + while time.time() < deadline: |
| 154 | + try: |
| 155 | + admin = AdminClient( |
| 156 | + {"bootstrap.servers": self.bootstrap_servers} |
| 157 | + ) |
| 158 | + admin.list_topics(timeout=2) |
| 159 | + logger.info("Kafka is ready.") |
| 160 | + return |
| 161 | + except Exception: |
| 162 | + time.sleep(1) |
| 163 | + raise TimeoutError( |
| 164 | + f"Kafka did not become ready within {timeout}s. " |
| 165 | + "Check Docker logs for details." |
| 166 | + ) |
| 167 | + |
| 168 | + # ------------------------------------------------------------------ |
| 169 | + # Topic management |
| 170 | + # ------------------------------------------------------------------ |
| 171 | + |
| 172 | + def create_topic( |
| 173 | + self, |
| 174 | + topic_name: str, |
| 175 | + num_partitions: int = 1, |
| 176 | + replication_factor: int = 1, |
| 177 | + ) -> None: |
| 178 | + """ |
| 179 | + Create a Kafka topic idempotently. |
| 180 | +
|
| 181 | + If the topic already exists the call succeeds silently. |
| 182 | + """ |
| 183 | + admin = AdminClient({"bootstrap.servers": self.bootstrap_servers}) |
| 184 | + new_topic = NewTopic( |
| 185 | + topic_name, |
| 186 | + num_partitions=num_partitions, |
| 187 | + replication_factor=replication_factor, |
| 188 | + ) |
| 189 | + futures = admin.create_topics([new_topic], operation_timeout=5) |
| 190 | + for topic, future in futures.items(): |
| 191 | + try: |
| 192 | + future.result() |
| 193 | + logger.info("Created topic '%s'.", topic) |
| 194 | + except Exception as exc: |
| 195 | + if "TOPIC_ALREADY_EXISTS" in str(exc): |
| 196 | + logger.debug("Topic '%s' already exists; skipping.", topic) |
| 197 | + else: |
| 198 | + raise |
| 199 | + |
| 200 | + def create_topics_if_not_exist( |
| 201 | + self, topic_names: List[str], num_partitions: int = 1 |
| 202 | + ) -> None: |
| 203 | + """Batch-create topics idempotently.""" |
| 204 | + for topic in topic_names: |
| 205 | + self.create_topic(topic, num_partitions=num_partitions) |
| 206 | + |
| 207 | + def clear_all_topics(self) -> None: |
| 208 | + """Delete every non-internal topic (fast data reset between tests).""" |
| 209 | + admin = AdminClient({"bootstrap.servers": self.bootstrap_servers}) |
| 210 | + try: |
| 211 | + metadata = admin.list_topics(timeout=5) |
| 212 | + to_delete = [ |
| 213 | + t for t in metadata.topics if not t.startswith("__") |
| 214 | + ] |
| 215 | + if to_delete: |
| 216 | + logger.debug("Deleting leftover topics: %s", to_delete) |
| 217 | + futures = admin.delete_topics(to_delete, operation_timeout=5) |
| 218 | + for _topic, fut in futures.items(): |
| 219 | + fut.result() |
| 220 | + except Exception as exc: |
| 221 | + logger.warning("Topic cleanup failed: %s", exc) |
0 commit comments