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" MODEL_ARTIFACTS = ( ( "construction-ppe-yolov8", 22_537_898, "31ef3ca04a17cf545f3fcfc64c4af8993a41d52ccc460e82aff01d5354603533", ), ( "ppe-6classes-yolov8n", 5_625_014, "07172ef3ae9e256c40a1fb0ce3eefe5547d90170645aa73dded0fffc382cdb31", ), ) 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 @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()