FabledCurator/tests/test_gpu_triage.py

"""Failure triage (#125): probe errored jobs' files, flag verdicts, recover.

The probe is the arbiter: reason strings only bucket the overview. A file that
passes checksum+decode is 'file_ok' (operational failure); anything else is a
'defect' — surfaced for recovery and excluded from /retry_errors.
"""
import hashlib

import pytest
from PIL import Image as PILImage
from sqlalchemy import select

from backend.app.models import (
    Artist,
    ExternalLink,
    GpuJob,
    ImageProvenance,
    ImageRecord,
    Post,
    Source,
)
from backend.app.services.ml.gpu_triage import (
    classify_reason,
    recover_defective_image,
    triage_errored_jobs,
)

pytestmark = pytest.mark.integration


def test_classify_reason_buckets():
    assert classify_reason(
        "no frames sampled from video — moov atom not found"
    ) == "truncated_or_corrupt"
    assert classify_reason("ffmpeg timed out after 1200s") == "timeout"
    assert classify_reason(
        "gave up after repeated transient failures: HTTPConnectionPool read timed out"
    ) == "transient"
    assert classify_reason(
        "poisoned: 10+ lease attempts without ever completing"
    ) == "poisoned"
    assert classify_reason("cannot identify image file") == "decode"
    assert classify_reason("something novel") == "other"
    assert classify_reason(None) == "other"


async def _errored_image(db, tmp_path, *, name, sha, content: bytes | None,
                         error="no frames sampled from video — moov atom not found"):
    """An ImageRecord (file written iff content is not None) + an errored job."""
    path = tmp_path / name
    if content is not None:
        path.write_bytes(content)
    img = ImageRecord(
        path=str(path), sha256=sha, size_bytes=1, mime="image/png",
        width=1, height=1, origin="imported_filesystem",
        integrity_status="unknown",
    )
    db.add(img)
    await db.flush()
    db.add(GpuJob(image_record_id=img.id, task="ccip", status="error",
                  error=error, attempts=3))
    await db.flush()
    return img


@pytest.mark.asyncio
async def test_triage_probes_and_splits_defect_vs_file_ok(db, tmp_path):
    # Healthy: a real PNG whose recorded sha matches its bytes → file_ok.
    ok_path = tmp_path / "fine.png"
    PILImage.new("RGB", (4, 4), (200, 30, 30)).save(ok_path)
    ok_sha = hashlib.sha256(ok_path.read_bytes()).hexdigest()
    ok = await _errored_image(db, tmp_path, name="fine.png", sha=ok_sha,
                              content=None, error="ffmpeg timed out after 1200s")
    # Corrupt: bytes don't match the recorded sha → defect.
    bad = await _errored_image(db, tmp_path, name="bad.png", sha="0" * 64,
                               content=b"not a real png")
    # Missing: no file on disk at all → defect (failed_verification).
    gone = await _errored_image(db, tmp_path, name="gone.png", sha="1" * 64,
                                content=None)
    await db.commit()

    summary = await db.run_sync(lambda s: triage_errored_jobs(s))
    assert summary["probed"] == 3
    assert summary["defect"] == 2
    assert summary["file_ok"] == 1
    assert summary["partial"] is False

    # Column selects, not ORM refresh — the sweep wrote via Core DML.
    rows = dict((await db.execute(
        select(GpuJob.image_record_id, GpuJob.triage_status)
        .where(GpuJob.status == "error")
    )).all())
    assert rows[ok.id] == "file_ok"
    assert rows[bad.id] == "defect"
    assert rows[gone.id] == "defect"
    verdicts = dict((await db.execute(
        select(ImageRecord.id, ImageRecord.integrity_status)
        .where(ImageRecord.id.in_([ok.id, bad.id, gone.id]))
    )).all())
    assert verdicts[ok.id] == "ok"
    assert verdicts[bad.id] == "corrupt"
    assert verdicts[gone.id] == "failed_verification"

    # Idempotent: everything already triaged → no re-probe.
    again = await db.run_sync(lambda s: triage_errored_jobs(s))
    assert again["probed"] == 0


@pytest.mark.asyncio
async def test_recover_without_pollable_source_reports_no_source(db, tmp_path):
    img = await _errored_image(db, tmp_path, name="orphan.png", sha="2" * 64,
                               content=b"x")
    await db.commit()
    res = await db.run_sync(
        lambda s: recover_defective_image(s, img.id, images_root=tmp_path)
    )
    assert res["status"] == "no_source"
    still_there = (await db.execute(
        select(ImageRecord.id).where(ImageRecord.id == img.id)
    )).scalar_one_or_none()
    assert still_there == img.id


@pytest.mark.asyncio
async def test_recover_deletes_record_and_requeues_source(
    client, db, tmp_path, monkeypatch,
):
    img = await _errored_image(db, tmp_path, name="fixme.png", sha="3" * 64,
                               content=b"x")
    artist = Artist(name="Recov", slug="recov")
    db.add(artist)
    await db.flush()
    src = Source(artist_id=artist.id, platform="patreon",
                 url="https://www.patreon.com/recov", enabled=True)
    db.add(src)
    await db.flush()
    post = Post(artist_id=artist.id, source_id=src.id, external_post_id="p1")
    db.add(post)
    await db.flush()
    db.add(ImageProvenance(image_record_id=img.id, post_id=post.id,
                           source_id=src.id))
    # A settled external-host link on the same post: recovery must reset it
    # (the SURGICAL path — deep posts are never re-walked by the cadence).
    link = ExternalLink(post_id=post.id, artist_id=artist.id, host="mega",
                        url="https://mega.nz/file/x#key", status="downloaded",
                        attempts=2)
    db.add(link)
    await db.flush()
    img_id, src_id, link_id = img.id, src.id, link.id
    await db.commit()

    queued = []
    monkeypatch.setattr(
        "backend.app.tasks.download.download_source.delay",
        lambda sid: queued.append(sid),
    )
    fetches = []
    monkeypatch.setattr(
        "backend.app.tasks.external.fetch_external_link.delay",
        lambda lid: fetches.append(lid),
    )

    resp = await client.post(f"/api/gpu/errors/{img_id}/recover")
    assert resp.status_code == 200
    body = await resp.get_json()
    assert body["status"] == "refetch_queued"
    assert body["source_id"] == src_id
    assert body["links_reset"] == 1
    assert queued == [src_id]
    assert fetches == [link_id]

    # The link is armed for a fresh attempt.
    row = (await db.execute(
        select(ExternalLink.status, ExternalLink.attempts)
        .where(ExternalLink.id == link_id)
    )).one()
    assert tuple(row) == ("pending", 0)

    # Record gone — the error tombstones cascade away with it.
    remaining = (await db.execute(
        select(ImageRecord.id).where(ImageRecord.id == img_id)
    )).scalar_one_or_none()
    assert remaining is None
    jobs_left = (await db.execute(
        select(GpuJob.id).where(GpuJob.image_record_id == img_id)
    )).scalars().all()
    assert jobs_left == []