nat-engine 1__py3-none-any.whl

mannf/product/integrations/kafka_adapter.py

@@ -0,0 +1,342 @@
+# Copyright (C) 2026 Brad Guider
+# This file is part of NAT (Neural Agent Testing Framework).
+# Licensed under the AGPL-3.0. See LICENSE for details.
+# Commercial licensing available — see COMMERCIAL_LICENSE.md.
+
+"""Kafka consumer/producer test adapter.
+
+:class:`KafkaAdapter` provides a test adapter for Apache Kafka topics that
+implements the same ``execute()`` / ``health_check()`` interface as
+:class:`~mannf.product.integrations.http_sut.HttpApiSUT`.
+
+Supports:
+
+* **Producer tests** — publish a message to a topic and verify no errors.
+* **Consumer tests** — subscribe to a topic and verify a message arrives
+  within a configurable timeout.
+* **Round-trip tests** — produce a message, then consume from a reply topic
+  to verify end-to-end delivery.
+
+The adapter uses ``kafka-python`` when available.  When not installed, all
+execute() calls return a skip result so NAT continues functioning.
+"""
+
+from __future__ import annotations
+
+import json
+import logging
+import time
+import uuid
+from typing import Any, Dict, List, Optional
+
+from mannf.core.distributed.endpoint import Endpoint, EndpointRegistry
+from mannf.core.distributed.system_under_test import SystemUnderTest
+from mannf.core.testing.models import TestCase, TestResult
+
+logger = logging.getLogger(__name__)
+
+
+def _kafka_available() -> bool:
+    """Return True if the ``kafka-python`` package is installed."""
+    try:
+        import kafka  # noqa: F401
+        return True
+    except ImportError:
+        return False
+
+
+def _bytes_passthrough(value: bytes) -> bytes:
+    """Identity serializer for Kafka producers — pass bytes through unchanged."""
+    return value
+
+
+class KafkaAdapter(SystemUnderTest):
+    """Execute Kafka topic tests (produce / consume / round-trip).
+
+    Parameters
+    ----------
+    bootstrap_servers:
+        Kafka broker address(es) (e.g. ``"localhost:9092"`` or a list).
+    timeout:
+        Per-operation timeout in seconds (default: 10).
+    consumer_group:
+        Consumer group ID (default: ``"nat-test-group-<uuid>"``)
+        to avoid interfering with production consumers.
+    producer_config:
+        Additional ``kafka-python`` producer configuration overrides.
+    consumer_config:
+        Additional ``kafka-python`` consumer configuration overrides.
+    """
+
+    def __init__(
+        self,
+        bootstrap_servers: str | list[str] = "localhost:9092",
+        timeout: float = 10.0,
+        consumer_group: Optional[str] = None,
+        producer_config: Optional[Dict[str, Any]] = None,
+        consumer_config: Optional[Dict[str, Any]] = None,
+    ) -> None:
+        if isinstance(bootstrap_servers, str):
+            self.bootstrap_servers = [bootstrap_servers]
+        else:
+            self.bootstrap_servers = list(bootstrap_servers)
+
+        self.timeout = timeout
+        self.consumer_group = consumer_group or f"nat-test-group-{uuid.uuid4().hex[:8]}"
+        self.producer_config: Dict[str, Any] = producer_config or {}
+        self.consumer_config: Dict[str, Any] = consumer_config or {}
+
+        host = self.bootstrap_servers[0].split(":")[0]
+        try:
+            port = int(self.bootstrap_servers[0].split(":")[1])
+        except (IndexError, ValueError):
+            port = 9092
+
+        self._registry = EndpointRegistry()
+        self._registry.register(
+            Endpoint(
+                name="default",
+                host=host,
+                port=port,
+                protocol="kafka",
+            )
+        )
+
+    # ------------------------------------------------------------------
+    # SystemUnderTest interface
+    # ------------------------------------------------------------------
+
+    @property
+    def registry(self) -> EndpointRegistry:
+        return self._registry
+
+    async def execute(self, test_case: TestCase) -> TestResult:
+        """Execute a Kafka test described by *test_case*.
+
+        ``test_case.target`` is the Kafka **topic** name.
+
+        ``test_case.inputs`` may contain:
+
+        * ``action`` — one of ``"produce"`` (default), ``"consume"``,
+          ``"round_trip"``.
+        * ``payload`` — dict or string to publish (for produce/round_trip).
+        * ``reply_topic`` — topic to consume from in round_trip mode.
+        * ``max_poll_records`` — max records to consume (default: 1).
+        """
+        if not _kafka_available():
+            return TestResult(
+                test_case_id=test_case.id,
+                passed=False,
+                actual_output=None,
+                execution_time_ms=0.0,
+                error=(
+                    "kafka-python package is not installed. "
+                    "Install it with: pip install kafka-python"
+                ),
+            )
+
+        action: str = str(test_case.inputs.get("action", "produce")).lower()
+        topic: str = test_case.target
+
+        if action == "produce":
+            return await self._action_produce(test_case, topic)
+        elif action == "consume":
+            return await self._action_consume(test_case, topic)
+        elif action == "round_trip":
+            return await self._action_round_trip(test_case, topic)
+        else:
+            return TestResult(
+                test_case_id=test_case.id,
+                passed=False,
+                actual_output=None,
+                execution_time_ms=0.0,
+                error=f"Unknown Kafka action: {action!r}",
+            )
+
+    async def health_check(self) -> bool:
+        """Check that the Kafka broker is reachable."""
+        if not _kafka_available():
+            return False
+        try:
+            from kafka import KafkaAdminClient
+
+            admin = KafkaAdminClient(
+                bootstrap_servers=self.bootstrap_servers,
+                request_timeout_ms=int(min(self.timeout, 5.0) * 1000),
+            )
+            admin.close()
+            return True
+        except Exception:  # noqa: BLE001
+            return False
+
+    # ------------------------------------------------------------------
+    # Action helpers
+    # ------------------------------------------------------------------
+
+    async def _action_produce(self, test_case: TestCase, topic: str) -> TestResult:
+        """Publish a message to *topic*."""
+        from kafka import KafkaProducer
+        from kafka.errors import KafkaError
+
+        payload = test_case.inputs.get("payload") or {}
+        if isinstance(payload, dict):
+            value_bytes = json.dumps(payload).encode("utf-8")
+        else:
+            value_bytes = str(payload).encode("utf-8")
+
+        start = time.monotonic()
+        try:
+            producer = KafkaProducer(
+                bootstrap_servers=self.bootstrap_servers,
+                value_serializer=_bytes_passthrough,
+                request_timeout_ms=int(self.timeout * 1000),
+                **self.producer_config,
+            )
+            future = producer.send(topic, value=value_bytes)
+            record_metadata = future.get(timeout=self.timeout)
+            producer.close()
+            elapsed_ms = (time.monotonic() - start) * 1000.0
+            return TestResult(
+                test_case_id=test_case.id,
+                passed=True,
+                actual_output={
+                    "topic": record_metadata.topic,
+                    "partition": record_metadata.partition,
+                    "offset": record_metadata.offset,
+                },
+                execution_time_ms=elapsed_ms,
+                error=None,
+            )
+        except Exception as exc:  # noqa: BLE001
+            elapsed_ms = (time.monotonic() - start) * 1000.0
+            return TestResult(
+                test_case_id=test_case.id,
+                passed=False,
+                actual_output=None,
+                execution_time_ms=elapsed_ms,
+                error=f"Kafka produce error: {exc}",
+            )
+
+    async def _action_consume(self, test_case: TestCase, topic: str) -> TestResult:
+        """Consume at least one message from *topic* within the timeout."""
+        from kafka import KafkaConsumer
+
+        max_records = int(test_case.inputs.get("max_poll_records") or 1)
+        start = time.monotonic()
+        try:
+            consumer = KafkaConsumer(
+                topic,
+                bootstrap_servers=self.bootstrap_servers,
+                group_id=self.consumer_group,
+                auto_offset_reset="latest",
+                consumer_timeout_ms=int(self.timeout * 1000),
+                max_poll_records=max_records,
+                **self.consumer_config,
+            )
+            records: list[dict] = []
+            for msg in consumer:
+                try:
+                    value = json.loads(msg.value.decode("utf-8"))
+                except Exception:  # noqa: BLE001
+                    value = msg.value.decode("utf-8", errors="replace")
+                records.append({
+                    "topic": msg.topic,
+                    "partition": msg.partition,
+                    "offset": msg.offset,
+                    "value": value,
+                })
+                if len(records) >= max_records:
+                    break
+            consumer.close()
+            elapsed_ms = (time.monotonic() - start) * 1000.0
+
+            if not records:
+                return TestResult(
+                    test_case_id=test_case.id,
+                    passed=False,
+                    actual_output={"records": []},
+                    execution_time_ms=elapsed_ms,
+                    error="No messages received within timeout",
+                )
+
+            return TestResult(
+                test_case_id=test_case.id,
+                passed=True,
+                actual_output={"records": records, "count": len(records)},
+                execution_time_ms=elapsed_ms,
+                error=None,
+            )
+        except Exception as exc:  # noqa: BLE001
+            elapsed_ms = (time.monotonic() - start) * 1000.0
+            return TestResult(
+                test_case_id=test_case.id,
+                passed=False,
+                actual_output=None,
+                execution_time_ms=elapsed_ms,
+                error=f"Kafka consume error: {exc}",
+            )
+
+    async def _action_round_trip(self, test_case: TestCase, topic: str) -> TestResult:
+        """Produce to *topic* then consume from a reply topic."""
+        from kafka import KafkaConsumer, KafkaProducer
+
+        reply_topic = str(test_case.inputs.get("reply_topic") or f"{topic}.reply")
+        payload = test_case.inputs.get("payload") or {}
+
+        # Embed correlation ID for matching
+        correlation_id = uuid.uuid4().hex
+        if isinstance(payload, dict):
+            payload = {**payload, "_nat_correlation_id": correlation_id}
+        value_bytes = json.dumps(payload).encode("utf-8")
+
+        start = time.monotonic()
+        try:
+            producer = KafkaProducer(
+                bootstrap_servers=self.bootstrap_servers,
+                value_serializer=_bytes_passthrough,
+                request_timeout_ms=int(self.timeout * 1000),
+                **self.producer_config,
+            )
+            future = producer.send(topic, value=value_bytes)
+            future.get(timeout=self.timeout)
+            producer.close()
+
+            # Now consume from reply topic
+            consumer = KafkaConsumer(
+                reply_topic,
+                bootstrap_servers=self.bootstrap_servers,
+                group_id=f"{self.consumer_group}-reply",
+                auto_offset_reset="latest",
+                consumer_timeout_ms=int(self.timeout * 1000),
+                max_poll_records=10,
+                **self.consumer_config,
+            )
+
+            reply_found = False
+            for msg in consumer:
+                try:
+                    val = json.loads(msg.value.decode("utf-8"))
+                    if isinstance(val, dict) and val.get("_nat_correlation_id") == correlation_id:
+                        reply_found = True
+                        break
+                except Exception:  # noqa: BLE001
+                    pass
+            consumer.close()
+
+            elapsed_ms = (time.monotonic() - start) * 1000.0
+            return TestResult(
+                test_case_id=test_case.id,
+                passed=reply_found,
+                actual_output={"round_trip_complete": reply_found},
+                execution_time_ms=elapsed_ms,
+                error=None if reply_found else "Reply message not received within timeout",
+            )
+        except Exception as exc:  # noqa: BLE001
+            elapsed_ms = (time.monotonic() - start) * 1000.0
+            return TestResult(
+                test_case_id=test_case.id,
+                passed=False,
+                actual_output=None,
+                execution_time_ms=elapsed_ms,
+                error=f"Kafka round-trip error: {exc}",
+            )