cmvr_edge_ai/tests/integration/test_model_artifacts.py

from __future__ import annotations

import hashlib
from pathlib import Path

import pytest

from cmvr_edge_ai.config import load_config


PROJECT_ROOT = Path(__file__).resolve().parents[2]
CONSTRUCTION_WEIGHTS = "models/detection/construction-ppe-yolov8/v1/best.pt"
PEOPLE_TALKING_WEIGHTS = "models/detection/people-talking-yolov8x/v1/best.pt"
MOBILE_PHONE_WEIGHTS = (
    "models/detection/yolov8n-mobile-phone/yolov8n-mobile-phone.pt"
)
MODEL_ARTIFACTS = (
    (
        "construction-ppe-yolov8",
        22_537_898,
        "31ef3ca04a17cf545f3fcfc64c4af8993a41d52ccc460e82aff01d5354603533",
    ),
    (
        "ppe-6classes-yolov8n",
        5_625_014,
        "07172ef3ae9e256c40a1fb0ce3eefe5547d90170645aa73dded0fffc382cdb31",
    ),
    (
        "people-talking-yolov8x",
        136_694_313,
        "86cd63926de1c69f70dd7f8755e4de9e382bbdeb0668821b7c3846b826471526",
    ),
)


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_name", "expected_size", "expected_sha256"),
    MODEL_ARTIFACTS,
)
def test_model_artifact_is_complete(
    model_name: str,
    expected_size: int,
    expected_sha256: str,
) -> None:
    model_dir = PROJECT_ROOT / "models" / "detection" / model_name / "v1"
    weights = model_dir / "best.pt"

    assert (model_dir / "README.md").is_file()
    assert weights.is_file()
    assert weights.stat().st_size == expected_size
    assert _sha256(weights) == expected_sha256


def test_mobile_phone_model_artifact_is_complete() -> None:
    weights = PROJECT_ROOT / MOBILE_PHONE_WEIGHTS
    model_card = weights.parent / "README.md"

    assert model_card.is_file()
    assert weights.is_file()
    assert weights.stat().st_size == 6_234_666
    assert (
        _sha256(weights)
        == "9230e4bfa7cba7134215c4c7f228b5e58760b02138788ccaff0149258c2d2e19"
    )


@pytest.mark.parametrize(
    ("config_path", "pipeline_id"),
    (
        (PROJECT_ROOT / "configs" / "edge_ai.yaml", "detection"),
        (
            PROJECT_ROOT / "configs" / "debug" / "detection_viewer.yaml",
            "detection_show",
        ),
    ),
)
def test_detection_configs_use_repository_construction_weights(
    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@1"
    assert detector.params["model_options"]["weights"] == CONSTRUCTION_WEIGHTS
    assert (PROJECT_ROOT / CONSTRUCTION_WEIGHTS).is_file()


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

    assert detector.params["model"] == "people-talking-yolov8x@1"
    assert detector.params["detect_labels"] == ["talking on phone"]
    assert detector.params["model_options"]["weights"] == PEOPLE_TALKING_WEIGHTS
    assert (PROJECT_ROOT / PEOPLE_TALKING_WEIGHTS).is_file()
    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
    )


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

    assert detector.params["model"] == "yolov8n-mobile-phone@1"
    assert detector.params["detect_labels"] == ["mobile_phone"]
    assert detector.params["model_options"]["weights"] == MOBILE_PHONE_WEIGHTS
    assert ppe_detector.params["attach_frame"] is True
    assert detector.params["attach_frame"] is True
    assert viewer.params["expected_inputs"] == ["ppe", "phone"]
    assert (PROJECT_ROOT / MOBILE_PHONE_WEIGHTS).is_file()
    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
    )