kafka source tests

Author: ovv

quixstreams/app.py

@@ -666,7 +666,7 @@ def get_producer(self) -> Producer:
         ```python
         from quixstreams import Application
 
-        app = Application.Quix(...)
+        app = Application(...)
         topic = app.topic("input")
 
         with app.get_producer() as producer:
@@ -706,7 +706,7 @@ def get_consumer(
         ```python
         from quixstreams import Application
 
-        app = Application.Quix(...)
+        app = Application(...)
         topic = app.topic("input")
 
         with app.get_consumer() as consumer:
@@ -719,6 +719,11 @@ def get_consumer(
                     # consumer.store_offsets(msg)
 
         ```
+
+        :param auto_commit_enable: Enable or disable auto commit
+            Default - True
+        :param consumer_group_suffix: Suffix added to the application consumer group
+            Default - ""
         """
         self._setup_topics()
 
@@ -731,6 +736,24 @@ def get_consumer(
         )
 
     def get_admin(self) -> TopicAdmin:
+        """
+        Create and return a TopicAdmin instance.
+        The TopicAdmin is initialized with params passed to Application
+
+        It's useful for configuring and inspecting kafka topics.
+
+        Example Snippet:
+
+        ```python
+        from quixstreams import Application
+
+        app = Application(...)
+        admin = app.get_admin()
+        topics = admin.list_topics():
+
+        print(topics)
+        ```
+        """
         return TopicAdmin(
             broker_address=self._broker_address,
             extra_config=self._producer_extra_config,

quixstreams/checkpointing/checkpoint.py

@@ -26,6 +26,15 @@
 
 
 class BaseCheckpoint:
+    """
+    Base class to keep track of state updates and consumer offsets and to checkpoint these
+    updates on schedule.
+
+    Two implementation exist:
+        * one for checkpointing the Application in quixstreams/checkpoint/checkpoint.py
+        * one for checkpointing the kafka source in quixstreams/sources/kafka/checkpoint.py
+    """
+
     def __init__(
         self,
         commit_interval: float,
@@ -106,8 +115,7 @@ def commit(self):
 
 class Checkpoint(BaseCheckpoint):
     """
-    Class to keep track of state updates and consumer offsets and to checkpoint these
-    updates on schedule.
+    Checkpoint implementation used by the application
     """
 
     def __init__(

quixstreams/sources/kafka/checkpoint.py

@@ -13,6 +13,10 @@
 
 
 class Checkpoint(BaseCheckpoint):
+    """
+    Checkpoint implementation used by the KafkaSource
+    """
+
     def __init__(
         self,
         producer: RowProducer,
@@ -33,6 +37,13 @@ def close(self):
         pass
 
     def commit(self):
+        """
+        Commit the checkpoint.
+
+        This method will:
+         2. Flush the producer to ensure everything is delivered.
+         3. Commit topic offsets.
+        """
         unproduced_msg_count = self._producer.flush(self._flush_timeout)
         if unproduced_msg_count > 0:
             raise CheckpointProducerTimeout(

quixstreams/sources/kafka/kafka.py

@@ -65,6 +65,7 @@ def __init__(
         commit_interval: float = 5,
         commit_every: int = 0,
         flush_timeout: float = 10,
+        admin_timeout: float = 10,
     ) -> None:
         super().__init__(name, shutdown_timeout)
 
@@ -100,6 +101,7 @@ def __init__(
         self._commit_interval = commit_interval
         self._commit_every = commit_every
         self._flush_timeout = flush_timeout
+        self._admin_timeout = admin_timeout
 
         self._running = True
         self._error: Optional[Exception] = None
@@ -135,7 +137,7 @@ def run(self) -> None:
             topic_name, partition, offset = msg.topic(), msg.partition(), msg.offset()
             self._checkpoint.store_offset(topic_name, partition, offset)
             logger.log(
-                5,
+                40,
                 "producing message to partition %s with offset %s",
                 partition,
                 offset,
@@ -177,13 +179,21 @@ def init_checkpoint(self) -> Checkpoint:
 
     def _validate_topics(self):
         source_topic = self._source_cluster_admin.inspect_topics(
-            topic_names=[self._topic], timeout=5
+            topic_names=[self._topic], timeout=self._admin_timeout
         ).get(self._topic)
+
+        if source_topic is None:
+            raise ValueError(f"Source topic {self._topic} not found")
+
         logger.debug("source topic %s configuration: %s", self._topic, source_topic)
 
         target_topic = self._target_cluster_admin.inspect_topics(
-            topic_names=[self._producer_topic.name], timeout=5
+            topic_names=[self._producer_topic.name], timeout=self._admin_timeout
         ).get(self._producer_topic.name)
+
+        if target_topic is None:
+            raise ValueError(f"Destination topic {self._producer_topic.name} not found")
+
         logger.debug(
             "destination topic %s configuration: %s",
             self._producer_topic.name,
@@ -203,7 +213,7 @@ def _target_cluster_offsets(self, partitions) -> Dict[int, int]:
             for partition in partitions
         ]
         partitions_commited = self._target_cluster_consumer.committed(
-            partitions, timeout=5
+            partitions, timeout=self._admin_timeout
         )
         return {
             partition.partition: partition.offset for partition in partitions_commited
@@ -248,7 +258,7 @@ def cleanup(self, failed):
 
     def default_topic(self) -> SourceTopic:
         config = self._source_cluster_admin.inspect_topics(
-            topic_names=[self._topic], timeout=5
+            topic_names=[self._topic], timeout=self._admin_timeout
         ).get(self._topic)
         if config is None:
             config = TopicConfig(num_partitions=1, replication_factor=1)

tests/conftest.py

@@ -34,14 +34,22 @@ def log_test_progress(request: pytest.FixtureRequest):
 
 
 @pytest.fixture(scope="session")
-def kafka_container() -> KafkaContainer:
-    (
-        kafka_container,
-        broker_address,
-        kafka_port,
-    ) = ContainerHelper.create_kafka_container()
-    test_logger.debug(f"Starting Kafka container on {broker_address}")
-    ContainerHelper.start_kafka_container(kafka_container)
-    test_logger.debug(f"Started Kafka container on {broker_address}")
-    yield KafkaContainer(broker_address=broker_address)
-    kafka_container.stop()
+def kafka_container_factory() -> KafkaContainer:
+    def factory():
+        (
+            kafka_container,
+            broker_address,
+            kafka_port,
+        ) = ContainerHelper.create_kafka_container()
+        test_logger.debug(f"Starting Kafka container on {broker_address}")
+        ContainerHelper.start_kafka_container(kafka_container)
+        test_logger.debug(f"Started Kafka container on {broker_address}")
+        yield KafkaContainer(broker_address=broker_address)
+        kafka_container.stop()
+
+    return factory
+
+
+@pytest.fixture(scope="session")
+def kafka_container(kafka_container_factory) -> KafkaContainer:
+    yield from kafka_container_factory()