cmvr_edge_ai/tests/unit/test_inference_contracts.py

from __future__ import annotations

import base64

import pytest
from pydantic import ValidationError

from cmvr_edge_ai.contracts import (
    CoordinateSpace,
    DetectionsOutput,
    ImageArtifact,
    ImageInput,
    InferenceError,
    InferenceRequest,
    InferenceResponse,
    InferenceStatus,
    ScalarOutput,
)


def _encoded(value: bytes = b"encoded-image") -> str:
    return base64.b64encode(value).decode("ascii")


def test_inference_request_round_trips_discriminated_image_input() -> None:
    request = InferenceRequest(
        request_id="request-1",
        category="detect.analog_gauge",
        source_id="robot-01/front-camera",
        captured_at_ns=123456789,
        inputs=(
            ImageInput(
                media_type="IMAGE/JPEG",
                data=_encoded(),
                width=1280,
                height=720,
            ),
        ),
        parameters={"include_stage_details": False},
        requested_artifact_roles=("annotated",),
    )

    wire = request.model_dump(mode="json")
    parsed = InferenceRequest.model_validate(wire)

    assert wire["schema_version"] == "cmvr.inference-request/v1"
    assert wire["inputs"][0]["kind"] == "image"
    assert wire["inputs"][0]["media_type"] == "image/jpeg"
    assert parsed == request
    assert isinstance(parsed.inputs[0], ImageInput)
    assert parsed.source_id == "robot-01/front-camera"
    assert parsed.captured_at_ns == 123456789


@pytest.mark.parametrize("data", ["", "not-base64", "YWJjZA==="])
def test_image_input_rejects_invalid_base64(data: str) -> None:
    with pytest.raises(ValidationError, match="base64|at least 1 character"):
        ImageInput(media_type="image/jpeg", data=data)


def test_inference_request_rejects_duplicate_input_names_and_unknown_fields() -> None:
    image = {
        "kind": "image",
        "name": "frame",
        "media_type": "image/jpeg",
        "encoding": "base64",
        "data": _encoded(),
    }
    with pytest.raises(ValidationError, match="input names must be unique"):
        InferenceRequest.model_validate(
            {
                "category": "detect.mobile_phone",
                "inputs": [image, image],
            }
        )

    with pytest.raises(ValidationError, match="extra_forbidden"):
        InferenceRequest.model_validate(
            {
                "category": "detect.mobile_phone",
                "inputs": [image],
                "model_path": "/tmp/arbitrary.pt",
            }
        )


def test_inference_response_parses_heterogeneous_outputs_and_image_artifact() -> None:
    response = InferenceResponse.model_validate(
        {
            "request_id": "request-1",
            "trace_id": "trace-1",
            "category": "detect.analog_gauge",
            "source_id": "robot-01/front-camera",
            "model": {
                "model_id": "analog-gauge-reader@1",
                "backend": "isolated-worker",
            },
            "status": "succeeded",
            "outputs": [
                {
                    "kind": "scalar",
                    "name": "reading",
                    "value": 4.2,
                    "unit": "bar",
                    "confidence": 0.87,
                },
                {
                    "kind": "detections",
                    "name": "gauges",
                    "coordinate_space": "pixel_xyxy",
                    "items": [
                        {
                            "label": "gauge",
                            "confidence": 0.96,
                            "box": {
                                "x_min": 120,
                                "y_min": 80,
                                "x_max": 560,
                                "y_max": 520,
                            },
                        }
                    ],
                },
            ],
            "artifacts": [
                {
                    "kind": "image",
                    "artifact_id": "annotated-1",
                    "role": "annotated",
                    "media_type": "image/jpeg",
                    "encoding": "base64",
                    "data": _encoded(b"annotated"),
                    "width": 1280,
                    "height": 720,
                }
            ],
        }
    )

    assert response.status is InferenceStatus.SUCCEEDED
    assert isinstance(response.outputs[0], ScalarOutput)
    assert isinstance(response.outputs[1], DetectionsOutput)
    assert isinstance(response.artifacts[0], ImageArtifact)
    assert response.source_id == "robot-01/front-camera"
    wire = response.model_dump(mode="json")
    assert wire["outputs"][0]["kind"] == "scalar"
    assert wire["artifacts"][0]["kind"] == "image"


def test_response_status_requires_consistent_error_contract() -> None:
    common = {
        "request_id": "request-1",
        "category": "detect.analog_gauge",
    }
    with pytest.raises(ValidationError, match="requires error"):
        InferenceResponse(**common, status="failed")

    with pytest.raises(ValidationError, match="must not contain error"):
        InferenceResponse(
            **common,
            status="succeeded",
            error=InferenceError(code="backend_failed", message="worker exited"),
        )

    failed = InferenceResponse(
        **common,
        status="failed",
        error=InferenceError(
            code="backend_unavailable",
            message="gauge worker is not ready",
            retryable=True,
            stage="routing",
        ),
    )
    assert failed.error is not None
    assert failed.error.retryable is True


def test_normalized_detections_reject_coordinates_above_one() -> None:
    with pytest.raises(ValidationError, match="normalized detection"):
        DetectionsOutput.model_validate(
            {
                "coordinate_space": CoordinateSpace.NORMALIZED_XYXY,
                "items": [
                    {
                        "label": "phone",
                        "confidence": 0.9,
                        "box": {
                            "x_min": 0.1,
                            "y_min": 0.1,
                            "x_max": 640,
                            "y_max": 0.8,
                        },
                    }
                ],
            }
        )


def test_wire_contracts_reject_non_json_parameter_values() -> None:
    with pytest.raises(ValidationError):
        InferenceRequest(
            category="detect.mobile_phone",
            inputs=(ImageInput(media_type="image/jpeg", data=_encoded()),),
            parameters={"callback": object()},  # type: ignore[dict-item]
        )