geointel/backend/tests/test_analysis_job_queue.py

"""Tiled GPU inference must not run inside an HTTP request.

A configured YOLO run walks up to ``YOLO_MAX_TILES`` tiles through the GPU.
Doing that in the request handler holds a worker thread for minutes, gives the
operator no progress, and times the client out before the result exists. The
run is queued as a Job instead and executed by a background worker, which is
the same pattern the AOI operations already use.
"""

from __future__ import annotations

from uuid import uuid4

import pytest

from app.core.errors import AppError
from app.models import AnalysisRun, Detection, Job
from app.services.analysis_job_worker import AnalysisJobWorker
from app.services.detection_service import DetectionService


class FakeQuery:
    def __init__(self, rows):
        self.rows = list(rows)

    def filter(self, *criteria):
        return self

    def update(self, values, **_kwargs) -> int:
        """Stand in for the conditional claim: succeeds while still queued."""

        claimed = 0
        for row in self.rows:
            if getattr(row, "status", None) == "queued":
                row.status = "running"
                claimed += 1
        return claimed

    def order_by(self, *_args):
        return self

    def limit(self, count):
        self.rows = self.rows[:count]
        return self

    def all(self):
        return list(self.rows)


class FakeSession:
    def __init__(self, objects=None, query_rows=None):
        self.objects = dict(objects or {})
        self.query_rows = query_rows or {}
        self.added = []

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

    def query(self, model):
        return FakeQuery(self.query_rows.get(model, []))

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

    def commit(self):
        return None

    def rollback(self):
        return None

    def refresh(self, _item):
        return None

    def close(self):
        return None


def _queued_job(**parameters) -> Job:
    payload = {
        "project_id": str(uuid4()),
        "dataset_id": str(uuid4()),
        "model_id": "yolo-configured",
        "confidence_threshold": 0.4,
        "class_filter": ["building"],
        "tile_manifest_path": "/tiles/manifest.json",
        "parameters_json": {},
    }
    payload.update(parameters)
    return Job(
        id=uuid4(),
        job_type="detection.run",
        status="queued",
        project_id=uuid4(),
        parameters_json=payload,
    )


def test_queued_detection_job_is_dispatched_to_the_detection_service(monkeypatch) -> None:
    job = _queued_job()
    db = FakeSession(objects={(Job, job.id): job}, query_rows={Job: [job]})
    calls: list[dict] = []

    def fake_run(**kwargs):
        calls.append(kwargs)
        return type(
            "Result",
            (),
            {
                "status": "success",
                "detection_count": 3,
                "analysis_run_id": uuid4(),
                "job_id": kwargs["existing_job"].id,
                "model_dump": lambda self, **_: {"status": "success", "detection_count": 3},
            },
        )()

    monkeypatch.setattr(DetectionService, "run_detection", staticmethod(fake_run))

    processed = AnalysisJobWorker.run_once(db=db)

    assert processed == 1
    assert calls[0]["model_id"] == "yolo-configured"
    assert calls[0]["confidence_threshold"] == 0.4
    assert calls[0]["tile_manifest_path"] == "/tiles/manifest.json"
    assert calls[0]["existing_job"] is job
    assert job.status == "success"


def test_a_failing_run_marks_the_job_failed_instead_of_leaving_it_running(monkeypatch) -> None:
    job = _queued_job()
    db = FakeSession(objects={(Job, job.id): job}, query_rows={Job: [job]})

    def exploding(**_kwargs):
        raise AppError(code="DETECTION_TILE_NOT_FOUND", message="missing tile", status_code=422)

    monkeypatch.setattr(DetectionService, "run_detection", staticmethod(exploding))

    processed = AnalysisJobWorker.run_once(db=db)

    assert processed == 1
    assert job.status == "failed"
    assert job.error_message == "missing tile"
    assert job.result_json["error_code"] == "DETECTION_TILE_NOT_FOUND"


def test_an_unexpected_error_still_closes_the_job(monkeypatch) -> None:
    job = _queued_job()
    db = FakeSession(objects={(Job, job.id): job}, query_rows={Job: [job]})

    def exploding(**_kwargs):
        raise RuntimeError("CUDA out of memory")

    monkeypatch.setattr(DetectionService, "run_detection", staticmethod(exploding))

    AnalysisJobWorker.run_once(db=db)

    assert job.status == "failed"
    assert job.result_json["error_code"] == "ANALYSIS_JOB_INTERNAL_ERROR"


def test_job_types_the_worker_does_not_own_are_left_alone() -> None:
    job = _queued_job()
    job.job_type = "raster.clip"
    db = FakeSession(objects={(Job, job.id): job}, query_rows={Job: [job]})

    assert AnalysisJobWorker.run_once(db=db) == 0
    assert job.status == "queued"


def test_enqueue_validates_before_accepting_the_job() -> None:
    """A bad request is rejected up front, not minutes later in the worker."""

    db = FakeSession()

    with pytest.raises(AppError) as exc_info:
        DetectionService.enqueue_detection(
            db=db,
            project_id=uuid4(),
            dataset_id=uuid4(),
            model_id="yolo-configured",
            confidence_threshold=0.4,
        )

    assert exc_info.value.code == "PROJECT_NOT_FOUND"