cmvr_edge_ai/tests/unit/test_server_components.py

from __future__ import annotations

import asyncio
import base64
from io import BytesIO

import pytest
from PIL import Image

from cmvr_edge_ai.contracts import BoundingBox, Detection, DetectionResult, ImageFrame
from cmvr_edge_ai.contracts.inference import (
    ImageInput,
    InferenceRequest,
    InferenceResponse,
    InferenceStatus,
)
from cmvr_edge_ai.core import ComponentContext
from cmvr_edge_ai.core.envelope import Envelope
from cmvr_edge_ai.server import InvocationBroker
from cmvr_edge_ai.server.components import (
    DetectionInferenceResponseOperator,
    InferenceImageDecodeOperator,
    InvocationRequestSource,
    InvocationResponseSink,
)


def _jpeg_bytes(color: tuple[int, int, int] = (255, 0, 0)) -> bytes:
    image = Image.new("RGB", (4, 3), color)
    output = BytesIO()
    image.save(output, format="JPEG", quality=95)
    return output.getvalue()


def _request(
    image: bytes,
    *,
    request_id: str = "request-1",
    category: str = "detect.mobile_phone",
    roles: tuple[str, ...] = (),
) -> InferenceRequest:
    return InferenceRequest(
        request_id=request_id,
        category=category,
        inputs=(
            ImageInput(
                media_type="image/jpeg",
                data=base64.b64encode(image).decode("ascii"),
            ),
        ),
        requested_artifact_roles=roles,
    )


def _context(broker: InvocationBroker[object, object]) -> ComponentContext:
    return ComponentContext(
        pipeline_id="remote",
        node_id="boundary",
        shutdown_event=asyncio.Event(),
        metadata={"invocation_broker": broker},
    )


def test_request_source_consumes_only_its_registered_category() -> None:
    async def exercise() -> None:
        broker: InvocationBroker[InferenceRequest, InferenceResponse] = (
            InvocationBroker()
        )
        broker.register_route("detect.mobile_phone", capacity=1)
        source = InvocationRequestSource(
            "requests", {"category": "detect.mobile_phone"}
        )
        await source.setup(_context(broker))
        request = _request(_jpeg_bytes())
        submitted = await broker.submit(
            request.category, request.request_id or "", request
        )

        stream = source.messages()
        emission = await anext(stream)
        assert emission.port == "requests"
        assert emission.envelope.payload is request
        assert emission.envelope.trace_id == "request-1"
        assert emission.envelope.attributes["invocation_category"] == request.category
        assert (
            emission.envelope.attributes["invocation_token"]
            == submitted.invocation_token
        )
        await broker.cancel("request-1")
        await broker.close()
        await stream.aclose()

    asyncio.run(exercise())


def test_image_decode_operator_produces_a_decoded_bgr_frame() -> None:
    async def exercise() -> ImageFrame:
        request = _request(_jpeg_bytes())
        operator = InferenceImageDecodeOperator("decode", {})
        output = await operator.process(
            Envelope(
                request,
                schema_name="InferenceRequest",
                attributes={
                    "invocation_request_id": "request-1",
                    "invocation_category": request.category,
                    "invocation_submitted_at_ns": 1,
                },
            )
        )
        return output.envelope.payload

    frame = asyncio.run(exercise())
    assert frame.width == 4
    assert frame.height == 3
    assert frame.pixel_format == "BGR8"
    assert frame.codec == "none"
    assert len(frame.data) == 4 * 3 * 3
    # JPEG is lossy, but red remains dominant in the BGR channel order.
    assert frame.data[2] > frame.data[1]
    assert frame.data[2] > frame.data[0]


def test_image_decode_rejects_media_type_mismatch() -> None:
    request = InferenceRequest(
        request_id="request-1",
        category="detect.mobile_phone",
        inputs=(
            ImageInput(
                media_type="image/png",
                data=base64.b64encode(_jpeg_bytes()).decode("ascii"),
            ),
        ),
    )

    with pytest.raises(ValueError, match="does not match decoded image type"):
        asyncio.run(
            InferenceImageDecodeOperator("decode", {}).process(
                Envelope(request, schema_name="InferenceRequest")
            )
        )


def test_detection_response_and_sink_complete_the_original_waiter() -> None:
    async def exercise() -> InferenceResponse:
        broker: InvocationBroker[InferenceRequest, InferenceResponse] = (
            InvocationBroker()
        )
        broker.register_route("detect.mobile_phone", capacity=1)
        request = _request(_jpeg_bytes(), roles=("annotated",))
        submitted = await broker.submit(
            request.category, request.request_id or "", request
        )
        await broker.receive(request.category)

        frame = ImageFrame(
            data=bytes((0, 0, 255)) * 12,
            width=4,
            height=3,
            pixel_format="BGR8",
        )
        result = DetectionResult(
            detections=(
                Detection(
                    label="mobile_phone",
                    confidence=0.9,
                    box=BoundingBox(0, 0, 3, 2),
                ),
            ),
            model_id="yolov8n-mobile-phone@1",
            model_name="YOLOv8n Mobile Phone",
            inference_ms=4.5,
            source_frame=frame,
        )
        response_operator = DetectionInferenceResponseOperator(
            "response",
            {
                "category": "detect.mobile_phone",
                "backend": "ultralytics-yolo",
            },
        )
        response_emission = await response_operator.process(
            Envelope(
                result,
                schema_name="DetectionResult",
                trace_id="request-1",
                attributes={
                    "invocation_request_id": "request-1",
                    "invocation_token": submitted.invocation_token,
                    "invocation_category": "detect.mobile_phone",
                    "invocation_submitted_at_ns": 1,
                    "requested_artifact_roles": ("annotated",),
                },
            )
        )
        sink = InvocationResponseSink("responses", {})
        await sink.setup(_context(broker))
        await sink.consume(response_emission.envelope)
        response = await broker.wait("request-1", timeout_s=0.1)
        await broker.close()
        return response

    response = asyncio.run(exercise())
    assert response.status is InferenceStatus.SUCCEEDED
    assert response.category == "detect.mobile_phone"
    assert response.model is not None
    assert response.model.model_id == "yolov8n-mobile-phone@1"
    assert response.outputs[0].kind == "detections"
    assert response.artifacts[0].role == "annotated"
    assert base64.b64decode(response.artifacts[0].data).startswith(b"\xff\xd8")


def test_response_sink_rejects_response_for_a_different_request_id() -> None:
    async def exercise() -> None:
        broker: InvocationBroker[InferenceRequest, InferenceResponse] = (
            InvocationBroker()
        )
        broker.register_route("detect.mobile_phone", capacity=1)
        request = _request(_jpeg_bytes())
        submitted = await broker.submit(
            request.category, request.request_id or "", request
        )
        await broker.receive(request.category)
        response = InferenceResponse(
            request_id="wrong-request",
            category=request.category,
            status=InferenceStatus.NO_RESULT,
        )
        envelope = Envelope(
            response,
            schema_name="InferenceResponse",
            attributes={
                "invocation_request_id": submitted.request_id,
                "invocation_token": submitted.invocation_token,
                "invocation_category": submitted.category,
            },
        )
        sink = InvocationResponseSink("responses", {})
        await sink.setup(_context(broker))

        with pytest.raises(ValueError, match="request_id does not match"):
            await sink.consume(envelope)

        assert await broker.cancel(submitted.request_id) is True
        await broker.close()

    asyncio.run(exercise())


def test_detection_response_uses_registry_backend_propagated_by_detector() -> None:
    result = DetectionResult(
        detections=(),
        model_id="yolov8n-mobile-phone@2",
        model_name="YOLOv8n Mobile Phone ONNX",
        inference_ms=1.0,
    )
    operator = DetectionInferenceResponseOperator(
        "response",
        {"category": "detect.mobile_phone"},
    )
    envelope = Envelope(
        result,
        schema_name="DetectionResult",
        trace_id="request-onnx",
        attributes={
            "invocation_request_id": "request-onnx",
            "invocation_category": "detect.mobile_phone",
            "invocation_submitted_at_ns": 1,
            "detection_model_backend": "onnxruntime-yolov8",
        },
    )

    emission = asyncio.run(operator.process(envelope))

    assert emission.envelope.payload.model is not None
    assert emission.envelope.payload.model.backend == "onnxruntime-yolov8"


def test_detection_response_rejects_stale_backend_override() -> None:
    result = DetectionResult(
        detections=(),
        model_id="yolov8n-mobile-phone@2",
        inference_ms=1.0,
    )
    operator = DetectionInferenceResponseOperator(
        "response",
        {
            "category": "detect.mobile_phone",
            "backend": "ultralytics-yolo",
        },
    )
    envelope = Envelope(
        result,
        schema_name="DetectionResult",
        attributes={
            "invocation_request_id": "request-onnx",
            "invocation_category": "detect.mobile_phone",
            "invocation_submitted_at_ns": 1,
            "detection_model_backend": "onnxruntime-yolov8",
        },
    )

    with pytest.raises(ValueError, match="does not match"):
        asyncio.run(operator.process(envelope))