from __future__ import annotations import asyncio import logging import threading from collections.abc import Awaitable, Callable, Mapping, Sequence from concurrent.futures import ThreadPoolExecutor from typing import Any, TypeVar import pytest from cmvr_edge_ai.contracts import BoundingBox, Detection, DetectionResult, ImageFrame from cmvr_edge_ai.core.component import ComponentContext, Emission from cmvr_edge_ai.core.envelope import Envelope from cmvr_edge_ai.detection.base import DetectionModel from cmvr_edge_ai.detection.operator import DetectionOperator from cmvr_edge_ai.detection.registry import ( DetectionModelRegistry, DetectionModelSpec, UnknownDetectionModelError, ) ResultT = TypeVar("ResultT") class RecordingDetectionModel(DetectionModel): def __init__(self, detections: Sequence[Detection] = ()) -> None: self.detections = tuple(detections) self.load_calls = 0 self.predict_calls = 0 self.close_calls = 0 self.predict_arguments: list[tuple[ImageFrame, tuple[str, ...], float]] = [] self.call_threads: dict[str, list[int]] = { "load": [], "predict": [], "close": [], } def load(self) -> None: self.load_calls += 1 self.call_threads["load"].append(threading.get_ident()) def predict( self, frame: ImageFrame, labels: tuple[str, ...], confidence: float, ) -> Sequence[Detection]: self.predict_calls += 1 self.predict_arguments.append((frame, labels, confidence)) self.call_threads["predict"].append(threading.get_ident()) return self.detections def close(self) -> None: self.close_calls += 1 self.call_threads["close"].append(threading.get_ident()) def _detection(label: str, confidence: float) -> Detection: return Detection( label=label, confidence=confidence, box=BoundingBox(x_min=1.0, y_min=2.0, x_max=11.0, y_max=22.0), ) def _frame(*, codec: str = "none") -> ImageFrame: return ImageFrame( data=bytes(range(12)), width=2, height=2, pixel_format="BGR8", codec=codec, is_key_frame=codec != "none", ) def _envelope(frame: ImageFrame | None = None) -> Envelope[ImageFrame]: return Envelope( payload=_frame() if frame is None else frame, schema_name="ImageFrame", schema_version=1, source_id="camera-right", sequence=42, captured_at_ns=1_234_000_000, received_at_ns=1_235_000_000, deadline_ns=1_500_000_000, trace_id="trace-detection-42", session_id="camera-session", attributes={"site": "park-a"}, ) def _registry( model: RecordingDetectionModel, *, labels: tuple[str, ...] = ("Worker", "No-Helmet", "No-Vest"), factory_options: list[dict[str, Any]] | None = None, ) -> DetectionModelRegistry: def factory(options: Mapping[str, Any]) -> DetectionModel: if factory_options is not None: factory_options.append(dict(options)) return model registry = DetectionModelRegistry() registry.register( DetectionModelSpec( model_id="construction-ppe@1", name="Construction PPE", supported_labels=labels, factory=factory, backend="fake", ) ) return registry def _context(executor: ThreadPoolExecutor) -> ComponentContext: return ComponentContext( pipeline_id="detection-test", node_id="detector", shutdown_event=asyncio.Event(), metadata={"thread_executor": executor}, ) async def _keep_restricted_event_loop_responsive() -> None: """Provide wakeups where sandboxed self-pipe notifications are unavailable.""" while True: await asyncio.sleep(0.01) def _run_with_thread_executor( exercise: Callable[[ThreadPoolExecutor], Awaitable[ResultT]], ) -> ResultT: async def runner() -> ResultT: loop = asyncio.get_running_loop() executor = ThreadPoolExecutor(max_workers=1) loop.set_default_executor(executor) ticker = asyncio.create_task(_keep_restricted_event_loop_responsive()) try: return await exercise(executor) finally: ticker.cancel() await asyncio.gather(ticker, return_exceptions=True) executor.shutdown(wait=True, cancel_futures=True) # asyncio.run() otherwise tries to shut down the same executor again. loop._default_executor = None # type: ignore[attr-defined] return asyncio.run(runner()) def test_operator_loads_once_predicts_and_closes_idempotently() -> None: model = RecordingDetectionModel([_detection("Worker", 0.9)]) options_seen: list[dict[str, Any]] = [] operator = DetectionOperator( "detector", { "model": "construction-ppe@1", "model_options": {"weights": "/models/ppe.pt", "device": "cpu"}, }, model_registry=_registry(model, factory_options=options_seen), ) async def scenario(executor: ThreadPoolExecutor) -> None: await operator.setup(_context(executor)) result = await operator.process(_envelope()) await operator.stop() await operator.stop() assert isinstance(result, Emission) assert model.load_calls == 1 assert model.predict_calls == 1 assert model.close_calls == 1 assert options_seen == [{"weights": "/models/ppe.pt", "device": "cpu"}] _run_with_thread_executor(scenario) def test_operator_filters_selected_labels_and_per_label_confidence() -> None: selected_worker = _detection("Worker", 0.50) selected_violation = _detection("No-Helmet", 0.80) model = RecordingDetectionModel( [ _detection("Worker", 0.49), selected_worker, _detection("No-Helmet", 0.79), selected_violation, _detection("No-Vest", 0.99), ] ) operator = DetectionOperator( "detector", { "model": "construction-ppe@1", "detect_labels": ["Worker", "No-Helmet"], "confidence": 0.50, "label_confidence": {"No-Helmet": 0.80}, }, model_registry=_registry(model), ) async def scenario(executor: ThreadPoolExecutor) -> None: await operator.setup(_context(executor)) emission = await operator.process(_envelope()) await operator.stop() assert isinstance(emission, Emission) payload = emission.envelope.payload assert isinstance(payload, DetectionResult) assert payload.detections == (selected_worker, selected_violation) assert operator.selected_labels == ("Worker", "No-Helmet") assert model.predict_arguments[0][1:] == (("Worker", "No-Helmet"), 0.50) _run_with_thread_executor(scenario) def test_operator_rejects_unknown_model_during_construction() -> None: with pytest.raises(UnknownDetectionModelError, match="missing@1"): DetectionOperator( "detector", {"model": "missing@1"}, model_registry=DetectionModelRegistry(), ) @pytest.mark.parametrize( "params", [ { "model": "construction-ppe@1", "detect_labels": ["Unknown-Label"], }, { "model": "construction-ppe@1", "label_confidence": {"Unknown-Label": 0.9}, }, ], ) def test_operator_rejects_unknown_labels_during_construction( params: dict[str, Any], ) -> None: with pytest.raises(ValueError, match="Unknown-Label"): DetectionOperator( "detector", params, model_registry=_registry(RecordingDetectionModel()), ) def test_operator_rejects_encoded_frame_without_calling_model() -> None: model = RecordingDetectionModel() operator = DetectionOperator( "detector", {"model": "construction-ppe@1"}, model_registry=_registry(model), ) async def scenario(executor: ThreadPoolExecutor) -> None: await operator.setup(_context(executor)) with pytest.raises(ValueError, match="decoded ImageFrame"): await operator.process(_envelope(_frame(codec="h264"))) await operator.stop() assert model.predict_calls == 0 _run_with_thread_executor(scenario) def test_operator_emits_result_metadata_and_preserves_envelope_correlation() -> None: source = _envelope() expected_detection = _detection("No-Helmet", 0.93) model = RecordingDetectionModel([expected_detection]) operator = DetectionOperator( "detector", {"model": "construction-ppe@1"}, model_registry=_registry(model), ) async def scenario(executor: ThreadPoolExecutor) -> None: await operator.setup(_context(executor)) emission = await operator.process(source) await operator.stop() assert isinstance(emission, Emission) assert emission.port == "detections" output = emission.envelope assert output.schema == "DetectionResult/v1" assert isinstance(output.payload, DetectionResult) assert output.payload.detections == (expected_detection,) assert output.payload.model_id == "construction-ppe@1" assert output.payload.model_name == "Construction PPE" assert output.payload.inference_ms >= 0.0 assert output.payload.source_frame is None assert operator.attach_frame is False assert output.source_id == source.source_id assert output.sequence == source.sequence assert output.captured_at_ns == source.captured_at_ns assert output.received_at_ns == source.received_at_ns assert output.deadline_ns == source.deadline_ns assert output.trace_id == source.trace_id assert output.session_id == source.session_id assert output.attributes == { "site": "park-a", "detection_model_id": "construction-ppe@1", "detection_model_name": "Construction PPE", } _run_with_thread_executor(scenario) def test_operator_attach_frame_keeps_the_decoded_frame_by_reference() -> None: source = _envelope() model = RecordingDetectionModel([_detection("No-Helmet", 0.93)]) operator = DetectionOperator( "detector", {"model": "construction-ppe@1", "attach_frame": True}, model_registry=_registry(model), ) async def scenario(executor: ThreadPoolExecutor) -> None: await operator.setup(_context(executor)) emission = await operator.process(source) await operator.stop() assert isinstance(emission, Emission) assert isinstance(emission.envelope.payload, DetectionResult) assert emission.envelope.payload.source_frame is source.payload assert operator.attach_frame is True _run_with_thread_executor(scenario) def test_blocking_model_methods_run_on_the_supplied_thread_pool() -> None: event_loop_thread = threading.get_ident() model = RecordingDetectionModel() operator = DetectionOperator( "detector", {"model": "construction-ppe@1"}, model_registry=_registry(model), ) async def scenario(executor: ThreadPoolExecutor) -> None: await operator.setup(_context(executor)) await operator.process(_envelope()) await operator.stop() assert all(len(thread_ids) == 1 for thread_ids in model.call_threads.values()) worker_threads = { thread_id for thread_ids in model.call_threads.values() for thread_id in thread_ids } assert len(worker_threads) == 1 assert event_loop_thread not in worker_threads _run_with_thread_executor(scenario) def test_max_fps_skips_frames_without_invoking_model_again() -> None: model = RecordingDetectionModel() operator = DetectionOperator( "detector", {"model": "construction-ppe@1", "max_fps": 0.001}, model_registry=_registry(model), ) async def scenario(executor: ThreadPoolExecutor) -> None: await operator.setup(_context(executor)) first = await operator.process(_envelope()) second = await operator.process(_envelope()) await operator.stop() assert isinstance(first, Emission) assert second is None assert model.predict_calls == 1 _run_with_thread_executor(scenario) def test_operator_logs_model_load_and_configured_inference_heartbeat( caplog: pytest.LogCaptureFixture, monkeypatch: pytest.MonkeyPatch, ) -> None: timestamps = iter((0, 1_000_000_000, 6_000_000_000)) monkeypatch.setattr( "cmvr_edge_ai.detection.operator.monotonic_ns", lambda: next(timestamps), ) model = RecordingDetectionModel([_detection("No-Helmet", 0.93)]) operator = DetectionOperator( "detector", { "model": "construction-ppe@1", "detect_labels": ["No-Helmet"], "inference_log_interval_s": 5, }, model_registry=_registry(model), ) async def scenario(executor: ThreadPoolExecutor) -> None: await operator.setup(_context(executor)) await operator.process(_envelope()) await operator.process(_envelope()) await operator.process(_envelope()) await operator.stop() with caplog.at_level(logging.INFO, logger="cmvr_edge_ai.detection.operator"): _run_with_thread_executor(scenario) messages = [record.getMessage() for record in caplog.records] assert any( "detection model loaded node=detector model=construction-ppe@1" in message for message in messages ) inference_messages = [ message for message in messages if message.startswith("detection inference ") ] assert len(inference_messages) == 2 assert "total_frames=1" in inference_messages[0] assert "window_frames=1" in inference_messages[0] assert "window_detections=1" in inference_messages[0] assert "hit_labels=No-Helmet:1" in inference_messages[0] assert "total_frames=3" in inference_messages[1] assert "window_frames=2" in inference_messages[1] assert "window_detections=2" in inference_messages[1] assert "hit_labels=No-Helmet:2" in inference_messages[1] @pytest.mark.parametrize("value", [0, -1, 3600.1, True]) def test_operator_rejects_invalid_inference_log_interval(value: Any) -> None: with pytest.raises(ValueError, match="inference_log_interval_s"): DetectionOperator( "detector", { "model": "construction-ppe@1", "inference_log_interval_s": value, }, model_registry=_registry(RecordingDetectionModel()), )