cmvr_edge_ai/tests/integration/test_model_artifacts.py

from __future__ import annotations

import hashlib
import json
from pathlib import Path

import pytest

from cmvr_edge_ai.application import create_default_model_registry
from cmvr_edge_ai.config import load_config
from cmvr_edge_ai.contracts import ImageFrame


PROJECT_ROOT = Path(__file__).resolve().parents[2]
CONSTRUCTION_ONNX = (
    "models/detection/construction-ppe-yolov8/v2/model.onnx"
)
PEOPLE_TALKING_ONNX = (
    "models/detection/people-talking-yolov8x/v2/model.onnx"
)
MOBILE_PHONE_ONNX = (
    "models/detection/yolov8n-mobile-phone/v2/model.onnx"
)

SOURCE_ARTIFACTS = (
    (
        "construction-ppe-yolov8@1",
        "models/detection/construction-ppe-yolov8/v1/best.pt",
        22_537_898,
        "31ef3ca04a17cf545f3fcfc64c4af8993a41d52ccc460e82aff01d5354603533",
    ),
    (
        "ppe-6classes-yolov8n@1",
        "models/detection/ppe-6classes-yolov8n/v1/best.pt",
        5_625_014,
        "07172ef3ae9e256c40a1fb0ce3eefe5547d90170645aa73dded0fffc382cdb31",
    ),
    (
        "people-talking-yolov8x@1",
        "models/detection/people-talking-yolov8x/v1/best.pt",
        136_694_313,
        "86cd63926de1c69f70dd7f8755e4de9e382bbdeb0668821b7c3846b826471526",
    ),
    (
        "yolov8n-mobile-phone@1",
        "models/detection/yolov8n-mobile-phone/yolov8n-mobile-phone.pt",
        6_234_666,
        "9230e4bfa7cba7134215c4c7f228b5e58760b02138788ccaff0149258c2d2e19",
    ),
)

ONNX_ARTIFACTS = (
    (
        "construction-ppe-yolov8@1",
        "construction-ppe-yolov8@2",
        CONSTRUCTION_ONNX,
        44_774_179,
        "b8e5d116a964d0e7091e14b5382ab9301c196ae33b363516df04fddfb0d0b57a",
    ),
    (
        "ppe-6classes-yolov8n@1",
        "ppe-6classes-yolov8n@2",
        "models/detection/ppe-6classes-yolov8n/v2/model.onnx",
        10_993_196,
        "6cb0e567b4fbf353aeb61c271ab33c45260d6ddddbe26eb3b64c245c2683c03b",
    ),
    (
        "people-talking-yolov8x@1",
        "people-talking-yolov8x@2",
        PEOPLE_TALKING_ONNX,
        272_787_930,
        "b4387307dbdf1083c0229a12e833ccdc4a3ad0063d5bdb7908765072470e840b",
    ),
    (
        "yolov8n-mobile-phone@1",
        "yolov8n-mobile-phone@2",
        MOBILE_PHONE_ONNX,
        12_265_264,
        "df170d9bc86984894797da6f2101519617669c3f589d75289acb0ba1937f857f",
    ),
)


def _sha256(path: Path) -> str:
    digest = hashlib.sha256()
    with path.open("rb") as artifact:
        for chunk in iter(lambda: artifact.read(1024 * 1024), b""):
            digest.update(chunk)
    return digest.hexdigest()


@pytest.mark.parametrize(
    ("model_id", "relative_path", "expected_size", "expected_sha256"),
    SOURCE_ARTIFACTS,
)
def test_pt_export_source_is_complete(
    model_id: str,
    relative_path: str,
    expected_size: int,
    expected_sha256: str,
) -> None:
    weights = PROJECT_ROOT / relative_path

    assert model_id.endswith("@1")
    assert (weights.parent / "README.md").is_file()
    assert weights.is_file()
    assert weights.stat().st_size == expected_size
    assert _sha256(weights) == expected_sha256


@pytest.mark.parametrize(
    (
        "source_model_id",
        "model_id",
        "relative_path",
        "expected_size",
        "expected_sha256",
    ),
    ONNX_ARTIFACTS,
)
def test_onnx_artifact_and_manifest_are_complete(
    source_model_id: str,
    model_id: str,
    relative_path: str,
    expected_size: int,
    expected_sha256: str,
) -> None:
    weights = PROJECT_ROOT / relative_path
    manifest_path = weights.with_name("manifest.json")
    manifest = json.loads(manifest_path.read_text(encoding="utf-8"))
    registry = create_default_model_registry(discover_entry_points=False)
    spec = registry.resolve(model_id)

    assert (weights.parent / "README.md").is_file()
    assert weights.is_file()
    assert weights.stat().st_size == expected_size
    assert _sha256(weights) == expected_sha256
    assert manifest["schema_version"] == "cmvr.detection-model-manifest/v1"
    assert manifest["source_model_id"] == source_model_id
    assert manifest["model_id"] == model_id
    assert manifest["backend"] == "onnxruntime-yolov8"
    assert manifest["artifact"] == relative_path
    assert manifest["artifact_sha256"] == expected_sha256
    assert manifest["labels"] == list(spec.supported_labels)
    assert manifest["input"]["shape"] == [1, 3, 640, 640]
    assert manifest["export"] == {
        "batch": 1,
        "dynamic": False,
        "format": "onnx",
        "half": False,
        "nms": False,
        "opset": 17,
        "simplify": True,
    }
    source = PROJECT_ROOT / manifest["source"]
    assert source.is_file()
    assert manifest["source_sha256"] == _sha256(source)


@pytest.mark.parametrize(
    ("model_id", "relative_path"),
    [(item[1], item[2]) for item in ONNX_ARTIFACTS],
)
def test_onnx_metadata_has_stable_identity_without_training_host_paths(
    model_id: str,
    relative_path: str,
) -> None:
    onnx = pytest.importorskip("onnx")
    weights = PROJECT_ROOT / relative_path
    model = onnx.load(str(weights), load_external_data=False)
    metadata = {item.key: item.value for item in model.metadata_props}
    manifest = json.loads(
        weights.with_name("manifest.json").read_text(encoding="utf-8")
    )

    assert "date" not in metadata
    assert "data.yaml" not in metadata["description"]
    assert metadata["cmvr_model_id"] == model_id
    assert metadata["cmvr_source_sha256"] == manifest["source_sha256"]


@pytest.mark.parametrize(
    "model_id",
    [item[1] for item in ONNX_ARTIFACTS],
)
def test_real_onnx_artifact_loads_and_runs_on_cpu(model_id: str) -> None:
    pytest.importorskip("onnxruntime")
    registry = create_default_model_registry(discover_entry_points=False)
    spec = registry.resolve(model_id)
    model = spec.factory(
        {
            "providers": ["CPUExecutionProvider"],
            "intra_op_threads": 1,
            "inter_op_threads": 1,
        }
    )
    frame = ImageFrame(
        data=bytes(32 * 32 * 3),
        width=32,
        height=32,
        pixel_format="BGR8",
    )

    try:
        model.load()
        detections = tuple(
            model.predict(frame, spec.supported_labels, confidence=0.99)
        )
    finally:
        model.close()

    for detection in detections:
        assert detection.label in spec.supported_labels
        assert 0.0 <= detection.box.x1 <= detection.box.x2 <= frame.width
        assert 0.0 <= detection.box.y1 <= detection.box.y2 <= frame.height


@pytest.mark.parametrize(
    ("config_path", "pipeline_id"),
    (
        (PROJECT_ROOT / "configs" / "active_detection.yaml", "detection"),
        (
            PROJECT_ROOT / "configs" / "detection_viewer.yaml",
            "detection_show",
        ),
        (PROJECT_ROOT / "configs" / "server_detect.yaml", "detect_ppe"),
    ),
)
def test_detection_configs_use_repository_construction_onnx(
    config_path: Path,
    pipeline_id: str,
) -> None:
    config = load_config(config_path)
    detector = config.pipelines[pipeline_id].nodes["detector"]

    assert detector.params["model"] == "construction-ppe-yolov8@2"
    assert detector.params["model_options"]["weights"] == CONSTRUCTION_ONNX
    assert detector.params["model_options"]["providers"] == [
        "CPUExecutionProvider"
    ]
    assert "device" not in detector.params["model_options"]
    assert "half" not in detector.params["model_options"]


def test_production_pipeline_uses_repository_people_talking_onnx() -> None:
    config = load_config(PROJECT_ROOT / "configs" / "active_detection.yaml")
    pipeline = config.pipelines["detection"]
    detector = pipeline.nodes["phone_detector"]
    gate = pipeline.nodes["phone_repeat_gate"]

    assert detector.params["model"] == "people-talking-yolov8x@2"
    assert detector.params["detect_labels"] == ["talking on phone"]
    assert detector.params["model_options"]["weights"] == PEOPLE_TALKING_ONNX
    assert gate.params["rules"][0]["labels"] == ["talking on phone"]
    assert any(
        edge.source == "decoder.frames" and edge.target == "phone_detector.frames"
        for edge in pipeline.edges
    )
    assert any(
        edge.source == "phone_detector.detections"
        and edge.target == "phone_repeat_gate.detections"
        for edge in pipeline.edges
    )
    assert any(
        edge.source == "phone_repeat_gate.alerts"
        and edge.target == "alert_platform.input"
        for edge in pipeline.edges
    )


@pytest.mark.parametrize(
    ("config_name", "pipeline_id"),
    (
        ("detection_viewer.yaml", "detection_show"),
        ("server_detect.yaml", "detect_mobile_phone"),
    ),
)
def test_mobile_phone_configs_use_repository_onnx(
    config_name: str,
    pipeline_id: str,
) -> None:
    config = load_config(PROJECT_ROOT / "configs" / config_name)
    detector = config.pipelines[pipeline_id].nodes["phone_detector" if pipeline_id == "detection_show" else "detector"]

    assert detector.params["model"] == "yolov8n-mobile-phone@2"
    assert detector.params["detect_labels"] == ["mobile_phone"]
    assert detector.params["model_options"]["weights"] == MOBILE_PHONE_ONNX


def test_debug_viewer_preserves_same_frame_pairing() -> None:
    config = load_config(PROJECT_ROOT / "configs" / "detection_viewer.yaml")
    pipeline = config.pipelines["detection_show"]
    ppe_detector = pipeline.nodes["detector"]
    detector = pipeline.nodes["phone_detector"]
    viewer = pipeline.nodes["viewer"]

    assert ppe_detector.params["attach_frame"] is True
    assert detector.params["attach_frame"] is True
    assert viewer.params["expected_inputs"] == ["ppe", "phone"]
    assert any(
        edge.source == "frame_fanout.output"
        and edge.target == "phone_detector.frames"
        for edge in pipeline.edges
    )
    assert any(
        edge.source == "phone_detector.detections"
        and edge.target == "viewer.phone"
        for edge in pipeline.edges
    )
    assert any(
        edge.source == "detector.detections" and edge.target == "viewer.ppe"
        for edge in pipeline.edges
    )