326 lines
12 KiB
Python
326 lines
12 KiB
Python
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"
|