geointel/backend/tests/test_sprint13_yolo_preflight.py

from __future__ import annotations

from hashlib import sha256
import json
import os
import subprocess
import sys
from pathlib import Path

from fastapi.testclient import TestClient

from app.core.config import Settings
from app.main import app
from app.services.runtime_model_provenance_service import RuntimeModelProvenanceService
from app.services.yolo_preflight_service import YoloPreflightService


ROOT = Path(__file__).resolve().parents[2]


class AvailableAdapter:
    @staticmethod
    def dependencies_available() -> bool:
        return True

    def __init__(self, settings: Settings) -> None:
        self.settings = settings

    def load_model(self, model_path: Path) -> object:
        return {"model_path": str(model_path)}


class MissingDependencyAdapter:
    @staticmethod
    def dependencies_available() -> bool:
        return False


class FailingLoadAdapter(AvailableAdapter):
    def load_model(self, model_path: Path) -> object:
        raise RuntimeError(f"cannot load {model_path}")


def _manifest(tmp_path: Path, tile_count: int = 1) -> Path:
    tiles = []
    for index in range(tile_count):
        tile_path = tmp_path / f"tile_{index:04d}.tif"
        tile_path.write_bytes(b"tile")
        tiles.append(
            {
                "path": str(tile_path),
                "pixel_window": [0, 0, 100, 100],
                "bounds": [4.0, 51.0, 5.0, 52.0],
                "transform": [4.0, 0.01, 0.0, 52.0, 0.0, -0.01],
                "index": index,
            }
        )
    manifest_path = tmp_path / "manifest.json"
    manifest_path.write_text(json.dumps({"tiles": tiles, "count": tile_count}), encoding="utf-8")
    return manifest_path


def _write_model_sidecar(model_path: Path, settings: Settings) -> None:
    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": settings.yolo_model_id,
            "task_type": "object_detection",
            "sha256": model_sha256,
            "model_format": "pytorch",
            "framework": "ultralytics/pytorch",
            "class_mapping": {"0": "building"},
            "source_version": settings.yolo_model_version or "test-v1",
        },
        "source": {
            "source_registry_id": "11111111-1111-4111-8111-111111111111",
            "source_snapshot_id": "22222222-2222-4222-8222-222222222222",
            "source_registry_key": "model",
            "source_snapshot_checksum_sha256": model_sha256,
        },
        "lineage": {
            "upstream_asset_ids": ["test-training-corpus"],
            "upstream_checksums_sha256": ["a" * 64],
            "transformations": [
                {"name": "test-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)
    RuntimeModelProvenanceService.manifest_path_for_model(model_path).write_text(
        json.dumps(payload, sort_keys=True),
        encoding="utf-8",
    )


def test_yolo_preflight_reports_disabled_without_loading_model(tmp_path: Path, monkeypatch) -> None:
    monkeypatch.setenv("YOLO_CONFIG_DIR", str(tmp_path / "ultralytics"))

    result = YoloPreflightService.run(
        settings=Settings(yolo_enabled=False, yolo_model_path=str(tmp_path / "missing.pt")),
        tile_manifest_path=str(tmp_path / "missing-manifest.json"),
        yolo_adapter_class=AvailableAdapter,
    )

    assert result["status"] == "not_configured"
    assert result["checks"]["enabled"] is False
    assert result["checks"]["dependencies_available"] is None
    assert result["checks"]["model_file_exists"] is None
    assert result["checks"]["manifest_valid"] is None
    assert result["runtime"]["dependencies_assumed"] is False
    assert result["runtime"]["model_directory"] == str(tmp_path)
    assert result["runtime"]["yolo_config_dir"] == str(tmp_path / "ultralytics")
    assert "torch_version" in result["runtime"]
    assert "ultralytics_version" in result["runtime"]
    assert "cuda_available" in result["runtime"]


def test_yolo_preflight_distinguishes_missing_dependencies_from_missing_model(tmp_path: Path) -> None:
    model_path = tmp_path / "model.pt"
    model_path.write_bytes(b"weights")

    result = YoloPreflightService.run(
        settings=Settings(yolo_enabled=True, yolo_model_path=str(model_path)),
        tile_manifest_path=str(_manifest(tmp_path)),
        yolo_adapter_class=MissingDependencyAdapter,
    )

    assert result["status"] == "dependency_unavailable"
    assert result["checks"]["dependencies_available"] is False
    assert result["checks"]["model_file_exists"] is None
    assert result["checks"]["manifest_valid"] is None


def test_yolo_preflight_validates_model_and_manifest_without_importing_yolo(tmp_path: Path) -> None:
    model_path = tmp_path / "model.pt"
    model_path.write_bytes(b"weights")
    manifest_path = _manifest(tmp_path, tile_count=2)
    settings = Settings(
        storage_root=str(tmp_path),
        yolo_enabled=True,
        yolo_model_path=str(model_path),
        yolo_max_tiles=4,
    )
    _write_model_sidecar(model_path, settings)

    result = YoloPreflightService.run(
        settings=settings,
        tile_manifest_path=str(manifest_path),
        yolo_adapter_class=AvailableAdapter,
    )

    assert result["status"] == "ready"
    assert result["checks"]["dependencies_available"] is True
    assert result["checks"]["model_file_exists"] is True
    assert result["checks"]["model_provenance_valid"] is True
    assert result["checks"]["manifest_valid"] is True
    assert result["tile_count"] == 2
    assert result["will_download_models"] is False
    assert result["will_run_inference"] is False
    assert result["runtime"]["dependencies_assumed"] is False


def test_yolo_preflight_marks_assumed_dependencies_in_runtime_details(tmp_path: Path) -> None:
    model_path = tmp_path / "model.pt"
    model_path.write_bytes(b"weights")
    manifest_path = _manifest(tmp_path, tile_count=1)
    settings = Settings(
        storage_root=str(tmp_path),
        yolo_enabled=True,
        yolo_model_path=str(model_path),
        yolo_max_tiles=4,
    )
    _write_model_sidecar(model_path, settings)

    result = YoloPreflightService.run(
        settings=settings,
        tile_manifest_path=str(manifest_path),
        yolo_adapter_class=MissingDependencyAdapter,
        assume_dependencies=True,
    )

    assert result["status"] == "ready"
    assert result["checks"]["dependencies_available"] is True
    assert result["runtime"]["dependencies_assumed"] is True
    assert result["runtime"]["model_directory"] == str(tmp_path)


def test_yolo_preflight_can_explicitly_smoke_load_local_model(tmp_path: Path) -> None:
    model_path = tmp_path / "model.pt"
    model_path.write_bytes(b"weights")
    manifest_path = _manifest(tmp_path)
    settings = Settings(
        storage_root=str(tmp_path),
        yolo_enabled=True,
        yolo_model_path=str(model_path),
        yolo_max_tiles=4,
    )
    _write_model_sidecar(model_path, settings)

    result = YoloPreflightService.run(
        settings=settings,
        tile_manifest_path=str(manifest_path),
        yolo_adapter_class=AvailableAdapter,
        check_model_load=True,
        allow_offline_model_load=True,
    )

    assert result["status"] == "ready"
    assert result["checks"]["model_load_requested"] is True
    assert result["checks"]["model_load_ok"] is True
    assert result["will_download_models"] is False
    assert result["will_run_inference"] is False


def test_yolo_preflight_reports_explicit_model_load_failure(tmp_path: Path) -> None:
    model_path = tmp_path / "model.pt"
    model_path.write_bytes(b"weights")
    settings = Settings(
        storage_root=str(tmp_path),
        yolo_enabled=True,
        yolo_model_path=str(model_path),
        yolo_max_tiles=4,
    )
    _write_model_sidecar(model_path, settings)

    result = YoloPreflightService.run(
        settings=settings,
        tile_manifest_path=str(_manifest(tmp_path)),
        yolo_adapter_class=FailingLoadAdapter,
        check_model_load=True,
        allow_offline_model_load=True,
    )

    assert result["status"] == "model_load_failed"
    assert result["error_code"] == "DETECTION_MODEL_LOAD_FAILED"
    assert result["checks"]["model_load_ok"] is False


def test_yolo_preflight_script_outputs_json(tmp_path: Path) -> None:
    model_path = tmp_path / "model.pt"
    model_path.write_bytes(b"weights")
    manifest_path = _manifest(tmp_path)
    _write_model_sidecar(model_path, Settings(yolo_enabled=True, yolo_model_path=str(model_path)))

    result = subprocess.run(
        [
            sys.executable,
            str(ROOT / "scripts" / "yolo_preflight.py"),
            "--model-path",
            str(model_path),
            "--tile-manifest-path",
            str(manifest_path),
            "--assume-dependencies",
            "--json",
        ],
        cwd=ROOT,
        check=True,
        capture_output=True,
        text=True,
        # The script reads process settings; the manifest it is asked to
        # validate lives here, so this is the storage root for that run.
        env={**os.environ, "STORAGE_ROOT": str(tmp_path)},
    )
    payload = json.loads(result.stdout)

    assert payload["status"] == "ready"
    assert payload["model_path"] == str(model_path)
    assert payload["tile_manifest_path"] == str(manifest_path)


def test_yolo_preflight_script_uses_environment_configuration(tmp_path: Path, monkeypatch) -> None:
    model_path = tmp_path / "model.pt"
    model_path.write_bytes(b"weights")
    manifest_path = _manifest(tmp_path)
    _write_model_sidecar(model_path, Settings(
        storage_root=str(tmp_path),
        yolo_enabled=True,
        yolo_model_path=str(model_path),
        yolo_max_tiles=4,
    ))
    monkeypatch.setenv("STORAGE_ROOT", str(tmp_path))
    monkeypatch.setenv("YOLO_ENABLED", "true")
    monkeypatch.setenv("YOLO_MODEL_PATH", str(model_path))
    monkeypatch.setenv("YOLO_MAX_TILES", "4")

    result = subprocess.run(
        [
            sys.executable,
            str(ROOT / "scripts" / "yolo_preflight.py"),
            "--tile-manifest-path",
            str(manifest_path),
            "--assume-dependencies",
            "--json",
        ],
        cwd=ROOT,
        check=True,
        capture_output=True,
        text=True,
    )
    payload = json.loads(result.stdout)

    assert payload["status"] == "ready"
    assert payload["checks"]["enabled"] is True
    assert payload["model_path"] == str(model_path)
    assert payload["max_tiles"] == 4


def test_yolo_preflight_refuses_unmanifested_local_weights(tmp_path: Path) -> None:
    model_path = tmp_path / "model.pt"
    model_path.write_bytes(b"unmanifested weights")

    result = YoloPreflightService.run(
        settings=Settings(yolo_enabled=True, yolo_model_path=str(model_path)),
        tile_manifest_path=str(_manifest(tmp_path)),
        yolo_adapter_class=AvailableAdapter,
    )

    assert result["status"] == "contract_incomplete"
    assert result["checks"]["model_provenance_valid"] is False
    assert result["error_code"] == "MODEL_PROVENANCE_MANIFEST_MISSING"


def test_yolo_preflight_script_rejects_assumed_dependencies_for_model_load(tmp_path: Path) -> None:
    result = subprocess.run(
        [
            sys.executable,
            str(ROOT / "scripts" / "yolo_preflight.py"),
            "--model-path",
            str(tmp_path / "model.pt"),
            "--assume-dependencies",
            "--check-model-load",
            "--json",
        ],
        cwd=ROOT,
        check=False,
        capture_output=True,
        text=True,
    )

    assert result.returncode != 0
    assert "--check-model-load cannot be combined with --assume-dependencies" in result.stderr


def test_yolo_preflight_api_returns_canonical_envelope(monkeypatch, tmp_path: Path) -> None:
    monkeypatch.setenv("STORAGE_ROOT", str(tmp_path))
    monkeypatch.setenv("YOLO_ENABLED", "false")
    monkeypatch.setenv("YOLO_MODEL_PATH", str(tmp_path / "missing.pt"))
    monkeypatch.setenv("YOLO_CONFIG_DIR", str(tmp_path / "ultralytics"))

    response = TestClient(app).get("/api/v1/detection/yolo/preflight")

    assert response.status_code == 200
    payload = response.json()
    assert set(payload) == {"data"}
    assert payload["data"]["status"] == "not_configured"
    assert payload["data"]["checks"]["enabled"] is False
    assert payload["data"]["runtime"]["model_directory"] == str(tmp_path)
    assert payload["data"]["runtime"]["yolo_config_dir"] == str(tmp_path / "ultralytics")
    assert payload["data"]["will_download_models"] is False
    assert payload["data"]["will_run_inference"] is False