geointel/backend/tests/test_model_asset_catalog.py

from __future__ import annotations

import json
from pathlib import Path
from uuid import uuid4

import pytest
from fastapi.testclient import TestClient

from app.core.config import Settings
from app.core.errors import AppError
from app.main import app
from app.models import AnalysisRun, Dataset, Detection, Job, Project
from app.services.detection_service import DetectionService
from app.services.model_asset_catalog_service import ModelAssetCatalogService


class FakeSession:
    def __init__(self, objects=None) -> None:
        self.objects = objects or {}
        self.added = []
        self.commits = 0
        self.refreshes = []

    def get(self, model, item_id):
        return self.objects.get((model, item_id))

    def add(self, item) -> None:
        self.added.append(item)
        if getattr(item, "id", None) is not None:
            self.objects[(item.__class__, item.id)] = item

    def commit(self) -> None:
        self.commits += 1

    def refresh(self, item) -> None:
        self.refreshes.append(item)


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

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

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

    def predict_tile(self, model, tile_path: Path, confidence_threshold: float) -> list[dict]:
        assert model["model_path"].endswith("building-detector.pt")
        return [
            {
                "class_name": "building",
                "confidence": 0.9,
                "bbox": [10.0, 20.0, 30.0, 40.0],
                "properties": {"adapter": "mock"},
            }
        ]


def _project_and_raster_dataset():
    project_id = uuid4()
    dataset_id = uuid4()
    project = Project(id=project_id, name="Geel")
    dataset = Dataset(
        id=dataset_id,
        project_id=project_id,
        name="source.tif",
        dataset_type="raster",
        source="user_upload",
        storage_path="storage/uploads/source.tif",
    )
    db = FakeSession(objects={(Project, project_id): project, (Dataset, dataset_id): dataset})
    return db, project_id, dataset_id


def _manifest(tmp_path: Path) -> Path:
    tile_path = tmp_path / "tile_0000.tif"
    tile_path.write_bytes(b"tile")
    manifest_path = tmp_path / "manifest.json"
    manifest_path.write_text(
        json.dumps(
            {
                "tile_set_id": "tiles-fixture",
                "count": 1,
                "tiles": [
                    {
                        "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": 0,
                    }
                ],
            }
        ),
        encoding="utf-8",
    )
    return manifest_path


def test_model_asset_catalog_lists_supported_local_model_files(tmp_path: Path) -> None:
    model_file = tmp_path / "building-detector.pt"
    model_file.write_bytes(b"local model")
    ignored_file = tmp_path / "notes.txt"
    ignored_file.write_text("ignore me", encoding="utf-8")
    settings = Settings(yolo_models_dir=str(tmp_path), yolo_model_path=str(model_file), yolo_enabled=True)

    response = ModelAssetCatalogService.list_assets(settings=settings)

    assert response.total == 1
    asset = response.items[0]
    assert asset.model_asset_id == "building-detector-pt"
    assert asset.filename == "building-detector.pt"
    assert asset.display_name == "building-detector"
    assert asset.model_path == str(model_file)
    assert asset.size_bytes == len(b"local model")
    assert len(asset.sha256) == 64
    assert asset.active is True
    assert asset.status == "approved"
    assert asset.will_download_models is False


def test_model_asset_catalog_resolves_known_asset(tmp_path: Path) -> None:
    model_file = tmp_path / "building-detector.pt"
    model_file.write_bytes(b"local model")
    settings = Settings(yolo_models_dir=str(tmp_path), yolo_enabled=True)

    asset = ModelAssetCatalogService.resolve_asset("building-detector-pt", settings=settings)

    assert asset.filename == "building-detector.pt"
    assert asset.model_path == str(model_file)


def test_model_asset_catalog_only_exposes_explicit_active_asset_in_runtime(tmp_path: Path) -> None:
    active_file = tmp_path / "approved-building-detector.pt"
    active_file.write_bytes(b"approved")
    (tmp_path / "training-smoke.pt").write_bytes(b"experiment")
    (tmp_path / "partial-checkpoint.pt").write_bytes(b"partial")
    settings = Settings(
        yolo_models_dir=str(tmp_path),
        yolo_model_path=str(active_file),
        yolo_enabled=True,
    )

    response = ModelAssetCatalogService.list_assets(settings=settings)

    assert response.total == 1
    assert response.items[0].filename == active_file.name
    assert response.items[0].active is True
    assert response.items[0].status == "approved"


def test_model_asset_catalog_rejects_unknown_asset(tmp_path: Path) -> None:
    settings = Settings(yolo_models_dir=str(tmp_path), yolo_enabled=True)

    with pytest.raises(AppError) as exc_info:
        ModelAssetCatalogService.resolve_asset("missing-model", settings=settings)

    assert exc_info.value.code == "DETECTION_MODEL_ASSET_NOT_FOUND"
    assert exc_info.value.status_code == 404


def test_model_assets_api_returns_canonical_envelope(monkeypatch, tmp_path: Path) -> None:
    model_file = tmp_path / "building-detector.pt"
    model_file.write_bytes(b"local model")
    monkeypatch.setenv("YOLO_MODELS_DIR", str(tmp_path))
    monkeypatch.setenv("YOLO_MODEL_PATH", str(model_file))

    response = TestClient(app).get("/api/v1/detection/model-assets")

    assert response.status_code == 200
    payload = response.json()
    assert set(payload) == {"data"}
    assert payload["data"]["total"] == 1
    assert payload["data"]["items"][0]["model_asset_id"] == "building-detector-pt"
    assert payload["data"]["items"][0]["active"] is True
    assert payload["data"]["items"][0]["will_download_models"] is False


def test_detection_run_persists_selected_model_asset_parameters(tmp_path: Path) -> None:
    model_file = tmp_path / "building-detector.pt"
    model_file.write_bytes(b"local model")
    db, project_id, dataset_id = _project_and_raster_dataset()
    settings = Settings(
        yolo_enabled=True,
        yolo_model_path=str(tmp_path / "default.pt"),
        yolo_models_dir=str(tmp_path),
        yolo_max_tiles=4,
    )

    result = DetectionService.run_detection(
        db=db,
        project_id=project_id,
        dataset_id=dataset_id,
        model_id="yolo-configured",
        model_asset_id="building-detector-pt",
        confidence_threshold=0.5,
        tile_manifest_path=str(_manifest(tmp_path)),
        settings=settings,
        yolo_adapter_class=MockYoloAdapter,
    )

    jobs = [item for item in db.added if isinstance(item, Job)]
    runs = [item for item in db.added if isinstance(item, AnalysisRun)]
    detections = [item for item in db.added if isinstance(item, Detection)]

    assert result.status == "success"
    assert result.detection_count == 1
    assert jobs[0].parameters_json["model_asset_id"] == "building-detector-pt"
    assert jobs[0].parameters_json["model_asset_path"] == str(model_file)
    assert len(jobs[0].parameters_json["model_asset_sha256"]) == 64
    assert runs[0].parameters_json["model_asset_id"] == "building-detector-pt"
    assert detections[0].model_name == "yolo-configured"