Refactored tests and scripts for better modularity and reusability

- Introduced a reusable Kafka container fixture for integration tests using Testcontainers.
- Modularized functions in `consumer.py` and `producer.py` for better testability.
- Refactored test cases to use shared Kafka fixture, reducing redundancy.
- Enhanced argument parsing in `consumer.py` for improved clarity and functionality.
- Simplified test logic for producer and consumer interaction.

Author: starichkov

consumer.py

@@ -11,6 +11,12 @@
 import argparse
 import json
 
+parser = argparse.ArgumentParser(description="Kafka Event Consumer")
+parser.add_argument("-e", "--event-type", help="Filter by event_type (optional)", required=False)
+parser.add_argument("-g", "--group-id", help="Kafka consumer group ID (optional)", required=False)
+parser.add_argument("-t", "--test-mode", action="store_true")
+args = parser.parse_args()
+
 """
 Apache Kafka Consumer Demo
 
@@ -43,39 +49,24 @@ def try_parse_json(value: bytes):
         return value.decode('utf-8', errors='replace')
 
 
-parser = argparse.ArgumentParser(description="Kafka Event Consumer")
-parser.add_argument("-e", "--event-type", help="Filter by event_type (optional)", required=False)
-parser.add_argument("-g", "--group-id", help="Kafka consumer group ID (optional)", required=False)
-parser.add_argument("-t", "--test-mode", action="store_true")
-args = parser.parse_args()
-
-
-def main():
-    bootstrap_servers = os.environ.get("KAFKA_BOOTSTRAP_SERVERS", "localhost:9092")
-    topic = os.environ.get("KAFKA_TOPIC", "test-topic")
-
-    # Initialize the Kafka consumer with configuration
-    # - topic: The Kafka topic to subscribe to (from environment variable)
-    # - bootstrap_servers: Connection string for the Kafka broker
-    # - auto_offset_reset='earliest': Start reading from the beginning of the topic if no committed offset exists
-    # - enable_auto_commit=True: Automatically commit offsets
-    consumer_args = {
-        'bootstrap_servers': bootstrap_servers,
-        'auto_offset_reset': 'earliest',
-        'enable_auto_commit': True
-    }
-
-    if args.group_id:
-        consumer_args['group_id'] = args.group_id
+def consume_events(topic, consumer_args, event_type=None, group_id=None):
+    """
+    Consume messages from a Kafka topic.
 
-    if args.test_mode:
-        consumer_args['consumer_timeout_ms'] = 3000  # pragma: no cover
+    This function creates a Kafka consumer with the provided arguments and listens
+    to a specified topic for incoming messages. Messages are parsed as JSON where
+    possible and logged appropriately. Filtering is available based on an optional
+    event type attribute in the message.
 
+    :param topic: Kafka topic to consume messages from.
+    :param consumer_args: Dictionary of arguments to configure the KafkaConsumer.
+    :param event_type: Optional. Filters messages by the `event_type` attribute if it's included in the message payload.
+    """
     consumer = KafkaConsumer(topic, **consumer_args)
 
     logger = get_logger("consumer")
 
-    logger.info(f"Polyglot consumer listening, consumer group: {args.group_id}\n")
+    logger.info(f"Polyglot consumer listening, consumer group: {group_id}\n")
 
     try:
         # Continuously poll for new messages
@@ -90,7 +81,7 @@ def main():
 
             # Display the message with an appropriate prefix based on its type
             if isinstance(parsed, dict):
-                if args.event_type and parsed.get("event_type") != args.event_type:
+                if event_type and parsed.get("event_type") != event_type:
                     continue  # Skip non-matching event
                 logger.info(
                     f"✅ JSON ({parsed['event_type']}) | key={key} | partition={partition} | offset={offset} → {parsed}")
@@ -104,5 +95,29 @@ def main():
         consumer.close()
 
 
+def main():
+    """Main function that runs when the script is executed directly"""
+    kafka_topic = os.environ.get("KAFKA_TOPIC", "test-topic")
+    kafka_bootstrap_servers = os.environ.get("KAFKA_BOOTSTRAP_SERVERS", "localhost:9092")
+
+    # Initialize the Kafka consumer with configuration
+    # - bootstrap_servers: Connection string for the Kafka broker
+    # - auto_offset_reset='earliest': Start reading from the beginning of the topic if no committed offset exists
+    # - enable_auto_commit=True: Automatically commit offsets
+    kafka_consumer_args = {
+        'bootstrap_servers': kafka_bootstrap_servers,
+        'auto_offset_reset': 'earliest',
+        'enable_auto_commit': True
+    }
+
+    if args.group_id:
+        kafka_consumer_args['group_id'] = args.group_id
+
+    if args.test_mode:
+        kafka_consumer_args['consumer_timeout_ms'] = 3000  # pragma: no cover
+
+    consume_events(kafka_topic, kafka_consumer_args, args.event_type, args.group_id)
+
+
 if __name__ == "__main__":
     main()

producer.py

@@ -29,10 +29,7 @@
 EVENT_TYPES = ["note_created", "note_updated", "note_deleted"]
 
 
-def main():
-    bootstrap_servers = os.environ.get("KAFKA_BOOTSTRAP_SERVERS", "localhost:9092")
-    topic = os.environ.get("KAFKA_TOPIC", "test-topic")
-
+def produce_events(bootstrap_servers, topic):
     # Initialize the Kafka producer with configuration
     # - bootstrap_servers: Connection string for the Kafka broker
     # - value_serializer: Function to convert Python objects to bytes
@@ -61,5 +58,13 @@ def main():
     producer.flush()
 
 
+def main():
+    """Main function that runs when the script is executed directly"""
+    kafka_bootstrap_servers = os.environ.get("KAFKA_BOOTSTRAP_SERVERS", "localhost:9092")
+    kafka_topic = os.environ.get("KAFKA_TOPIC", "test-topic")
+
+    produce_events(kafka_bootstrap_servers, kafka_topic)
+
+
 if __name__ == "__main__":
     main()

tests/test_cli_integration.py

@@ -3,77 +3,83 @@
 import time
 import os
 from kafka import KafkaProducer
+import pytest
 
 
-def test_producer_and_consumer_via_scripts(tmp_path):
-    topic = "test-topic"
-
+@pytest.fixture(scope="module")
+def kafka_container():
+    """Fixture that provides a reusable Kafka container for all tests."""
     with KafkaContainer(image="confluentinc/cp-kafka:7.9.2") as kafka:
-        bootstrap_servers = kafka.get_bootstrap_server()
+        yield kafka
 
-        env = os.environ.copy()
-        env["KAFKA_BOOTSTRAP_SERVERS"] = bootstrap_servers
-        env["KAFKA_TOPIC"] = topic
 
-        # 1. Run the producer
-        subprocess.run(
-            ["python", "producer.py"],
-            env=env,
-            check=True,
-        )
+def test_producer_and_consumer_via_scripts(tmp_path, kafka_container):
+    topic = "test-topic-producer-consumer-scripts"
 
-        time.sleep(1)
+    bootstrap_servers = kafka_container.get_bootstrap_server()
 
-        # 2. Capture consumer output
-        env["COVERAGE_PROCESS_START"] = ".coveragerc"  # Enable coverage for subprocess
+    env = os.environ.copy()
+    env["KAFKA_BOOTSTRAP_SERVERS"] = bootstrap_servers
+    env["KAFKA_TOPIC"] = topic
 
-        result = subprocess.run(
-            ["python", "consumer.py", "--test-mode"],
-            env=env,
-            check=True,
-            capture_output=True,
-            text=True,
-            timeout=10,
-        )
+    # 1. Run the producer
+    subprocess.run(
+        ["python", "producer.py"],
+        env=env,
+        check=True,
+    )
 
-        assert "note event" in result.stdout.lower()
+    time.sleep(1)
 
+    # 2. Capture consumer output
+    env["COVERAGE_PROCESS_START"] = ".coveragerc"  # Enable coverage for subprocess
 
-def test_plain_text_consumer(tmp_path):
-    topic = "test-topic"
+    result = subprocess.run(
+        ["python", "consumer.py", "--test-mode"],
+        env=env,
+        check=True,
+        capture_output=True,
+        text=True,
+        timeout=10,
+    )
 
-    with KafkaContainer(image="confluentinc/cp-kafka:7.9.2") as kafka:
-        bootstrap_servers = kafka.get_bootstrap_server()
-
-        # Create a producer that sends plain text messages
-        producer = KafkaProducer(
-            bootstrap_servers=bootstrap_servers,
-            key_serializer=lambda k: k.encode('utf-8') if k else None,
-            value_serializer=lambda v: v.encode('utf-8')  # Plain text serializer
-        )
-
-        # Send a plain text message
-        plain_text_message = "This is a plain text message"
-        producer.send(topic, key="plain-text-key", value=plain_text_message)
-        producer.flush()
-
-        time.sleep(1)
-
-        # Run the consumer and capture its output
-        env = os.environ.copy()
-        env["KAFKA_BOOTSTRAP_SERVERS"] = bootstrap_servers
-        env["KAFKA_TOPIC"] = topic
-        env["COVERAGE_PROCESS_START"] = ".coveragerc"  # Enable coverage for subprocess
-
-        result = subprocess.run(
-            ["python", "consumer.py", "--test-mode"],
-            env=env,
-            check=True,
-            capture_output=True,
-            text=True,
-            timeout=10,
-        )
-
-        # Verify that the plain text message was processed correctly
-        assert "📦 Plain" in result.stdout
-        assert plain_text_message in result.stdout
+    assert "note event" in result.stdout.lower()
+
+
+def test_plain_text_consumer(tmp_path, kafka_container):
+    topic = "test-topic-plain-text-consumer"
+
+    bootstrap_servers = kafka_container.get_bootstrap_server()
+
+    # Create a producer that sends plain text messages
+    producer = KafkaProducer(
+        bootstrap_servers=bootstrap_servers,
+        key_serializer=lambda k: k.encode('utf-8') if k else None,
+        value_serializer=lambda v: v.encode('utf-8')  # Plain text serializer
+    )
+
+    # Send a plain text message
+    plain_text_message = "This is a plain text message"
+    producer.send(topic, key="plain-text-key", value=plain_text_message)
+    producer.flush()
+
+    time.sleep(1)
+
+    # Run the consumer and capture its output
+    env = os.environ.copy()
+    env["KAFKA_BOOTSTRAP_SERVERS"] = bootstrap_servers
+    env["KAFKA_TOPIC"] = topic
+    env["COVERAGE_PROCESS_START"] = ".coveragerc"  # Enable coverage for subprocess
+
+    result = subprocess.run(
+        ["python", "consumer.py", "--test-mode"],
+        env=env,
+        check=True,
+        capture_output=True,
+        text=True,
+        timeout=10,
+    )
+
+    # Verify that the plain text message was processed correctly
+    assert "📦 Plain" in result.stdout
+    assert plain_text_message in result.stdout