cmvr_edge_ai/tests/unit/test_export_detection_onnx.py

from __future__ import annotations

import hashlib
import importlib.util
import json
import os
from pathlib import Path
from types import ModuleType, SimpleNamespace
from typing import Any

import pytest


PROJECT_ROOT = Path(__file__).resolve().parents[2]
SCRIPT = PROJECT_ROOT / "scripts" / "export_detection_onnx.py"


def _load_script() -> ModuleType:
    spec = importlib.util.spec_from_file_location("cmvr_export_detection_onnx", SCRIPT)
    assert spec is not None and spec.loader is not None
    module = importlib.util.module_from_spec(spec)
    # dataclasses resolves postponed annotations through sys.modules.
    import sys

    sys.modules[spec.name] = module
    spec.loader.exec_module(module)
    return module


exporter = _load_script()


def test_export_script_disables_ultralytics_autoinstall_and_has_four_models() -> None:
    assert os.environ["YOLO_AUTOINSTALL"] == "false"
    assert set(exporter.EXPORT_DEFINITIONS) == {
        "construction-ppe-yolov8@2",
        "ppe-6classes-yolov8n@2",
        "people-talking-yolov8x@2",
        "yolov8n-mobile-phone@2",
    }
    assert all(
        definition.output.as_posix().endswith("/v2/model.onnx")
        for definition in exporter.EXPORT_DEFINITIONS.values()
    )


def test_export_refuses_existing_destination_without_force(tmp_path: Path) -> None:
    source = tmp_path / "source.pt"
    output = tmp_path / "model.onnx"
    manifest = tmp_path / "manifest.json"
    source.write_bytes(b"trusted checkpoint")
    output.write_bytes(b"existing")

    with pytest.raises(FileExistsError, match="without --force"):
        exporter._validate_paths(
            source=source,
            output=output,
            manifest=manifest,
            force=False,
        )


def test_export_rejects_an_lfs_pointer_source(tmp_path: Path) -> None:
    source = tmp_path / "source.pt"
    source.write_text(
        "version https://git-lfs.github.com/spec/v1\n"
        "oid sha256:abc\nsize 123\n",
        encoding="utf-8",
    )

    with pytest.raises(ValueError, match="Git LFS pointer"):
        exporter._validate_paths(
            source=source,
            output=tmp_path / "model.onnx",
            manifest=tmp_path / "manifest.json",
            force=False,
        )


def test_fake_export_is_staged_validated_and_published_with_digest(
    tmp_path: Path,
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    definition = exporter.EXPORT_DEFINITIONS["yolov8n-mobile-phone@2"]
    source = tmp_path / "trusted.pt"
    output = tmp_path / "v2" / "model.onnx"
    manifest = output.parent / "manifest.json"
    source.write_bytes(b"trusted-checkpoint")
    export_calls: list[dict[str, Any]] = []
    runtime_paths: list[Path] = []

    class FakeYOLO:
        task = "detect"
        names = dict(enumerate(definition.labels))

        def __init__(self, path: str, *, task: str) -> None:
            assert task == "detect"
            self.path = Path(path)

        def export(self, **kwargs: Any) -> str:
            export_calls.append(kwargs)
            exported = self.path.with_suffix(".onnx")
            exported.write_bytes(b"validated-onnx")
            return str(exported)

    class FakeRuntime:
        def __init__(
            self,
            params: dict[str, Any],
            *,
            expected_labels: tuple[str, ...],
        ) -> None:
            assert expected_labels == definition.labels
            runtime_paths.append(Path(params["weights"]))

        def load(self) -> None:
            assert runtime_paths[-1].is_file()

        def close(self) -> None:
            return None

    fake_model = SimpleNamespace(
        metadata_props=[
            SimpleNamespace(
                key="description",
                value="trained on /private/training/data.yaml",
            ),
            SimpleNamespace(key="date", value="non-deterministic"),
        ]
    )
    sanitized_metadata: dict[str, str] = {}

    def set_model_props(model: Any, metadata: dict[str, str]) -> None:
        assert model is fake_model
        sanitized_metadata.update(metadata)

    fake_onnx = SimpleNamespace(
        __version__="test-onnx",
        load=lambda path: fake_model,
        helper=SimpleNamespace(set_model_props=set_model_props),
        save_model=lambda model, path: None,
        checker=SimpleNamespace(check_model=lambda model: None),
    )
    dependencies = {
        "YOLO": FakeYOLO,
        "onnx": fake_onnx,
        "onnxruntime": SimpleNamespace(__version__="test-ort"),
        "onnxslim": None,
        "torch": SimpleNamespace(__version__="test-torch"),
        "ultralytics": SimpleNamespace(__version__="test-ultralytics"),
    }
    monkeypatch.setattr(
        exporter,
        "_require_export_dependencies",
        lambda *, simplify: dependencies,
    )
    monkeypatch.setattr(exporter, "OnnxYoloModel", FakeRuntime)
    args = exporter.parse_args(
        [
            "--model-id",
            definition.target_model_id,
            "--source",
            str(source),
            "--output",
            str(output),
            "--manifest",
            str(manifest),
            "--imgsz",
            "32",
            "--opset",
            "17",
            "--no-simplify",
        ]
    )

    published_output, published_manifest = exporter.export_model(args)

    assert published_output == output
    assert published_manifest == manifest
    assert output.read_bytes() == b"validated-onnx"
    payload = json.loads(manifest.read_text(encoding="utf-8"))
    assert payload["model_id"] == definition.target_model_id
    assert payload["source_model_id"] == definition.source_model_id
    assert payload["artifact_sha256"] == hashlib.sha256(
        b"validated-onnx"
    ).hexdigest()
    assert payload["export"] == {
        "batch": 1,
        "dynamic": False,
        "format": "onnx",
        "half": False,
        "nms": False,
        "opset": 17,
        "simplify": False,
    }
    assert export_calls == [
        {
            "format": "onnx",
            "imgsz": 32,
            "batch": 1,
            "dynamic": False,
            "simplify": False,
            "opset": 17,
            "nms": False,
            "half": False,
            "device": "cpu",
        }
    ]
    assert runtime_paths and output.parent in runtime_paths[0].parents
    assert not list(output.parent.glob(".onnx-export-*"))
    assert "date" not in sanitized_metadata
    assert "/private/" not in sanitized_metadata["description"]
    assert sanitized_metadata["cmvr_model_id"] == definition.target_model_id
    assert sanitized_metadata["cmvr_source_sha256"] == hashlib.sha256(
        b"trusted-checkpoint"
    ).hexdigest()