from __future__ import annotations from hashlib import sha256 import json from pathlib import Path from uuid import uuid4 import pytest from app.core.errors import AppError from app.models import DatasetQuarantine, SourceRegistry, SourceSnapshot from app.services.runtime_model_provenance_service import RuntimeModelProvenanceService class FakeSession: """Explicit database double for production-runtime provenance tests.""" def __init__(self, objects: dict[tuple[type, object], object] | None = None) -> None: self.objects = objects or {} def get(self, model, item_id): return self.objects.get((model, item_id)) def _write_sidecar( model_path: Path, *, model_id: str = "yolo-configured", task_type: str = "object_detection", framework: str = "ultralytics/pytorch", source_version: str = "test-v1", source_registry_id: str | None = None, source_snapshot_id: str | None = None, ) -> Path: model_sha256 = sha256(model_path.read_bytes()).hexdigest() payload = { "schema_version": RuntimeModelProvenanceService.MANIFEST_SCHEMA_VERSION, "data_contract": { "key": "geointel.model.pytorch", "version": "1.0.0", }, "model": { "model_id": model_id, "task_type": task_type, "sha256": model_sha256, "model_format": "pytorch", "framework": framework, "class_mapping": {"0": "building"}, "source_version": source_version, }, "source": { "source_registry_id": source_registry_id or str(uuid4()), "source_snapshot_id": source_snapshot_id or str(uuid4()), "source_registry_key": "model", "source_snapshot_checksum_sha256": model_sha256, }, "lineage": { "upstream_asset_ids": ["training-corpus:test-v1"], "upstream_checksums_sha256": ["a" * 64], "transformations": [ { "name": "pytorch-training", "version": "1.0.0", "checksum_sha256": "b" * 64, } ], }, "metadata": { "training_manifest_sha256": "c" * 64, }, "imported_at": "2026-08-01T10:00:00+00:00", } payload["metadata"]["runtime_manifest_sha256"] = RuntimeModelProvenanceService.manifest_self_checksum(payload) sidecar_path = RuntimeModelProvenanceService.manifest_path_for_model(model_path) sidecar_path.write_text(json.dumps(payload, sort_keys=True), encoding="utf-8") return sidecar_path def _governed_model_database( *, source_registry_id, source_snapshot_id, model_checksum: str, source_version: str = "test-v1", ) -> tuple[FakeSession, SourceRegistry, SourceSnapshot]: registry = SourceRegistry( id=source_registry_id, source_key="model", display_name="Governed test model artifacts", classification="experimental", authority_name="GeoIntel test fixture", freshness_status="current", ingest_status="configured", ) snapshot = SourceSnapshot( id=source_snapshot_id, source_registry_id=source_registry_id, snapshot_key=f"model-{source_version}", source_version=source_version, checksum_sha256=model_checksum, freshness_status="current", ingest_status="ingested", ) return ( FakeSession( { (SourceRegistry, source_registry_id): registry, (SourceSnapshot, source_snapshot_id): snapshot, } ), registry, snapshot, ) def test_runtime_model_provenance_accepts_byte_bound_pytorch_sidecar_for_structural_preflight(tmp_path: Path) -> None: model_path = tmp_path / "model.pt" model_path.write_bytes(b"trusted local model bytes") sidecar_path = _write_sidecar(model_path, source_version="v1") evidence = RuntimeModelProvenanceService.validate_for_runtime( model_path=model_path, model_id="yolo-configured", task_type="object_detection", expected_model_version="v1", allowed_frameworks=("ultralytics/pytorch",), ) assert evidence.model_sha256 == sha256(model_path.read_bytes()).hexdigest() assert evidence.manifest_path == str(sidecar_path.resolve()) assert evidence.data_contract_key == "geointel.model.pytorch" assert evidence.data_contract_version == "1.0.0" assert len(evidence.validation_report_sha256) == 64 def test_production_runtime_requires_db_bound_model_source_snapshot(tmp_path: Path) -> None: model_path = tmp_path / "model.pt" model_path.write_bytes(b"governed local model bytes") source_registry_id = uuid4() source_snapshot_id = uuid4() _write_sidecar( model_path, source_version="v1", source_registry_id=str(source_registry_id), source_snapshot_id=str(source_snapshot_id), ) db, _, _ = _governed_model_database( source_registry_id=source_registry_id, source_snapshot_id=source_snapshot_id, model_checksum=sha256(model_path.read_bytes()).hexdigest(), source_version="v1", ) evidence = RuntimeModelProvenanceService.validate_for_production_runtime( db=db, model_path=model_path, model_id="yolo-configured", task_type="object_detection", expected_model_version="v1", allowed_frameworks=("ultralytics/pytorch",), ) assert evidence.source_registry_id == str(source_registry_id) assert evidence.source_snapshot_id == str(source_snapshot_id) assert evidence.source_snapshot_checksum_sha256 == evidence.model_sha256 def test_production_runtime_rejects_missing_database_source_binding(tmp_path: Path) -> None: model_path = tmp_path / "model.pt" model_path.write_bytes(b"unbound model bytes") _write_sidecar(model_path) with pytest.raises(AppError) as exc_info: RuntimeModelProvenanceService.validate_for_production_runtime( db=FakeSession(), model_path=model_path, model_id="yolo-configured", task_type="object_detection", ) assert exc_info.value.code == "MODEL_PROVENANCE_SOURCE_REGISTRY_NOT_FOUND" def test_production_runtime_requires_a_database_session(tmp_path: Path) -> None: model_path = tmp_path / "model.pt" model_path.write_bytes(b"model bytes") _write_sidecar(model_path) with pytest.raises(AppError) as exc_info: RuntimeModelProvenanceService.validate_for_production_runtime( db=None, model_path=model_path, model_id="yolo-configured", task_type="object_detection", ) assert exc_info.value.code == "MODEL_PROVENANCE_DATABASE_REQUIRED" @pytest.mark.parametrize( ("mutation", "expected_code"), ( ("registry_unsafe", "MODEL_PROVENANCE_SOURCE_REGISTRY_UNSAFE"), ("snapshot_registry_mismatch", "MODEL_PROVENANCE_SOURCE_SNAPSHOT_REGISTRY_MISMATCH"), ("snapshot_missing", "MODEL_PROVENANCE_SOURCE_SNAPSHOT_NOT_FOUND"), ("snapshot_quarantined", "MODEL_PROVENANCE_SOURCE_SNAPSHOT_UNSAFE"), ("snapshot_checksum_mismatch", "MODEL_PROVENANCE_DATABASE_SNAPSHOT_CHECKSUM_MISMATCH"), ("active_quarantine", "MODEL_PROVENANCE_SOURCE_SNAPSHOT_QUARANTINED"), ), ) def test_production_runtime_rejects_unsafe_or_inconsistent_database_snapshot( tmp_path: Path, mutation: str, expected_code: str, ) -> None: model_path = tmp_path / "model.pt" model_path.write_bytes(b"governed model bytes") source_registry_id = uuid4() source_snapshot_id = uuid4() _write_sidecar( model_path, source_registry_id=str(source_registry_id), source_snapshot_id=str(source_snapshot_id), ) db, registry, snapshot = _governed_model_database( source_registry_id=source_registry_id, source_snapshot_id=source_snapshot_id, model_checksum=sha256(model_path.read_bytes()).hexdigest(), ) if mutation == "registry_unsafe": registry.ingest_status = "quarantined" elif mutation == "snapshot_registry_mismatch": snapshot.source_registry_id = uuid4() elif mutation == "snapshot_missing": db.objects.pop((SourceSnapshot, source_snapshot_id)) elif mutation == "snapshot_quarantined": snapshot.ingest_status = "quarantined" elif mutation == "snapshot_checksum_mismatch": snapshot.checksum_sha256 = "f" * 64 elif mutation == "active_quarantine": snapshot.quarantines = [ DatasetQuarantine( source_snapshot_id=source_snapshot_id, stage="test", reason_code="test_active_quarantine", status="quarantined", ) ] with pytest.raises(AppError) as exc_info: RuntimeModelProvenanceService.validate_for_production_runtime( db=db, model_path=model_path, model_id="yolo-configured", task_type="object_detection", ) assert exc_info.value.code == expected_code def test_runtime_model_provenance_rejects_missing_sidecar(tmp_path: Path) -> None: model_path = tmp_path / "model.pt" model_path.write_bytes(b"unmanifested local model bytes") with pytest.raises(AppError) as exc_info: RuntimeModelProvenanceService.validate_for_runtime( model_path=model_path, model_id="yolo-configured", task_type="object_detection", ) assert exc_info.value.code == "MODEL_PROVENANCE_MANIFEST_MISSING" def test_runtime_model_provenance_rejects_model_bytes_tampered_after_manifest(tmp_path: Path) -> None: model_path = tmp_path / "model.pt" model_path.write_bytes(b"original local model bytes") _write_sidecar(model_path) model_path.write_bytes(b"tampered local model bytes") with pytest.raises(AppError) as exc_info: RuntimeModelProvenanceService.validate_for_runtime( model_path=model_path, model_id="yolo-configured", task_type="object_detection", ) assert exc_info.value.code == "MODEL_PROVENANCE_MODEL_CHECKSUM_MISMATCH" def test_runtime_model_provenance_rejects_tampered_manifest_contents(tmp_path: Path) -> None: model_path = tmp_path / "model.pt" model_path.write_bytes(b"original local model bytes") sidecar_path = _write_sidecar(model_path) payload = json.loads(sidecar_path.read_text(encoding="utf-8")) payload["model"]["class_mapping"]["1"] = "road" sidecar_path.write_text(json.dumps(payload, sort_keys=True), encoding="utf-8") with pytest.raises(AppError) as exc_info: RuntimeModelProvenanceService.validate_for_runtime( model_path=model_path, model_id="yolo-configured", task_type="object_detection", ) assert exc_info.value.code == "MODEL_PROVENANCE_MANIFEST_CHECKSUM_MISMATCH" def test_runtime_model_provenance_rejects_other_contract_even_if_structurally_valid(tmp_path: Path) -> None: model_path = tmp_path / "model.pt" model_path.write_bytes(b"local model bytes") sidecar_path = _write_sidecar(model_path) payload = json.loads(sidecar_path.read_text(encoding="utf-8")) payload["data_contract"]["key"] = "geointel.vector.geojson" payload["metadata"]["runtime_manifest_sha256"] = RuntimeModelProvenanceService.manifest_self_checksum(payload) sidecar_path.write_text(json.dumps(payload, sort_keys=True), encoding="utf-8") with pytest.raises(AppError) as exc_info: RuntimeModelProvenanceService.validate_for_runtime( model_path=model_path, model_id="yolo-configured", task_type="object_detection", ) assert exc_info.value.code == "MODEL_PROVENANCE_MANIFEST_INVALID"