cmvr_edge_ai/tests/integration/test_detection_pipeline.py

from __future__ import annotations

import asyncio
from collections.abc import Mapping, Sequence
from concurrent.futures import ThreadPoolExecutor
from io import BytesIO
from typing import Any

import pytest

from cmvr_edge_ai.application import create_default_registry
from cmvr_edge_ai.compiler import compile_pipeline
from cmvr_edge_ai.config import load_config_data
from cmvr_edge_ai.contracts import (
    BoundingBox,
    Detection,
    DetectionAlert,
    ImageFrame,
)
from cmvr_edge_ai.core import Envelope, Sink
from cmvr_edge_ai.detection import (
    DetectionModel,
    DetectionModelRegistry,
    DetectionModelSpec,
)
from cmvr_edge_ai.plugins import PluginKind, PluginSpec


class _AlwaysViolationModel(DetectionModel):
    def __init__(self) -> None:
        self.load_calls = 0
        self.predict_calls = 0
        self.close_calls = 0

    def load(self) -> None:
        self.load_calls += 1

    def predict(
        self,
        frame: ImageFrame,
        labels: tuple[str, ...],
        confidence: float,
    ) -> Sequence[Detection]:
        del frame, confidence
        self.predict_calls += 1
        if "No-Helmet" not in labels:
            return ()
        return (
            Detection(
                "No-Helmet",
                0.95,
                BoundingBox(1.0, 2.0, 20.0, 40.0),
            ),
        )

    def close(self) -> None:
        self.close_calls += 1


class _CollectAlertSink(Sink):
    def __init__(self) -> None:
        self.alerts: list[Envelope[Any]] = []

    async def consume(
        self, envelope: Envelope[Any], input_port: str = "alerts"
    ) -> None:
        del input_port
        self.alerts.append(envelope)


async def _keep_event_loop_responsive() -> None:
    while True:
        await asyncio.sleep(0.01)


def test_decoded_frames_flow_through_registered_model_and_repeat_gate() -> None:
    pillow_image = pytest.importorskip("PIL.Image")
    model = _AlwaysViolationModel()
    models = DetectionModelRegistry()
    models.register(
        DetectionModelSpec(
            model_id="fake-ppe@1",
            name="Fake PPE",
            supported_labels=("No-Helmet",),
            factory=lambda options: model,
            backend="fake",
        )
    )
    registry = create_default_registry(
        discover_entry_points=False,
        model_registry=models,
    )
    sink = _CollectAlertSink()

    def sink_factory(node_id: str, params: Mapping[str, Any]) -> Sink:
        del node_id, params
        return sink

    registry.register(
        PluginSpec(
            plugin_id="test.alert_sink@1",
            kind=PluginKind.SINK,
            factory=sink_factory,
            inputs={"alerts": "DetectionAlert/v1"},
        )
    )
    width, height = 64, 48
    frames = (
        ImageFrame(bytes((0, 0, 255)) * (width * height), width, height, "BGR8"),
        ImageFrame(bytes((0, 255, 0)) * (width * height), width, height, "BGR8"),
        ImageFrame(bytes((255, 0, 0)) * (width * height), width, height, "BGR8"),
    )
    config = load_config_data(
        {
            "api_version": "cmvr.edge.ai/v1",
            "pipelines": {
                "ppe": {
                    "nodes": {
                        "frames": {
                            "uses": "core.sequence_source@1",
                            "with": {
                                "items": list(frames),
                                "schema_name": "ImageFrame",
                            },
                        },
                        "detector": {
                            "uses": "detection.model@1",
                            "with": {
                                "model": "fake-ppe@1",
                                "detect_labels": ["No-Helmet"],
                                "attach_frame": True,
                            },
                        },
                        "gate": {
                            "uses": "detection.repeat_gate@1",
                            "with": {
                                "time_source": "received",
                                "alert_image": {
                                    "enabled": True,
                                    "jpeg_quality": 85,
                                },
                                "rules": [
                                    {
                                        "id": "no-helmet",
                                        "labels": ["No-Helmet"],
                                        "min_hits": 2,
                                        "window_ms": 10_000,
                                        "cooldown_ms": 30_000,
                                    }
                                ],
                            },
                        },
                        "sink": {"uses": "test.alert_sink@1"},
                    },
                    "edges": [
                        {"from": "frames.output", "to": "detector.frames"},
                        {
                            "from": "detector.detections",
                            "to": "gate.detections",
                        },
                        {"from": "gate.alerts", "to": "sink.alerts"},
                    ],
                }
            },
        }
    )

    gate_health_details: list[str] = []

    async def scenario() -> None:
        loop = asyncio.get_running_loop()
        executor = ThreadPoolExecutor(max_workers=2)
        loop.set_default_executor(executor)
        ticker = asyncio.create_task(_keep_event_loop_responsive())
        try:
            compiled = compile_pipeline(
                config,
                "ppe",
                registry,
                metadata={"thread_executor": executor},
            )
            await compiled.runtime.run()
            gate_health = await compiled.runtime.nodes["gate"].health()
            gate_health_details.append(gate_health.detail)
        finally:
            ticker.cancel()
            await asyncio.gather(ticker, return_exceptions=True)
            executor.shutdown(wait=True, cancel_futures=True)
            loop._default_executor = None  # type: ignore[attr-defined]

    asyncio.run(scenario())

    assert model.load_calls == 1
    assert model.predict_calls == 3
    assert model.close_calls == 1
    assert len(sink.alerts) == 1
    alert = sink.alerts[0].payload
    assert isinstance(alert, DetectionAlert)
    assert alert.rule_id == "no-helmet"
    assert alert.hit_count == 2
    assert alert.labels == ("No-Helmet",)
    assert sink.alerts[0].sequence == 1
    assert alert.image is not None
    assert alert.image.media_type == "image/jpeg"
    assert (alert.image.width, alert.image.height) == (width, height)
    assert alert.image.data.startswith(b"\xff\xd8")

    with pillow_image.open(BytesIO(alert.image.data)) as decoded:
        decoded.load()
        assert decoded.format == "JPEG"
        assert decoded.mode == "RGB"
        assert decoded.size == (width, height)
        # This pixel lies outside the annotation and proves that the evidence
        # image came from sequence 1 (green), the frame that reached min_hits.
        red, green, blue = decoded.getpixel((width - 1, height - 1))
        assert green > 200
        assert red < 40
        assert blue < 40

    assert gate_health_details
    assert "processed_frames=3" in gate_health_details[0]
    assert "alerts_emitted=1" in gate_health_details[0]
    assert "alert_images_encoded=1" in gate_health_details[0]