FabledCurator/tests/test_head_auto_apply.py

"""Earned auto-apply (#114). The sweep is numpy-only (no scikit-learn), so the
apply logic is tested directly via the sync session; the API guards (disabled /
dry-run / conflict) via the async client."""
import pytest
from sqlalchemy import select

from backend.app.models import (
    HeadAutoApplyRun,
    ImageRecord,
    MLSettings,
    Tag,
    TagHead,
    TagKind,
)
from backend.app.models.tag import image_tag
from backend.app.services.ml.heads import auto_apply_sweep

pytestmark = pytest.mark.integration


def _emb(slot: int) -> list[float]:
    v = [0.0] * 1152
    v[slot] = 3.0
    return v


def _img(db, sha: str, emb) -> ImageRecord:
    img = ImageRecord(
        path=f"/images/{sha}.jpg", sha256=sha, size_bytes=1, mime="image/jpeg",
        width=1, height=1, origin="imported_filesystem",
        integrity_status="unknown", siglip_embedding=emb,
    )
    db.add(img)
    db.flush()
    return img


def _head(db, tag_id: int, slot: int, *, threshold=0.5, n_pos=30):
    s = db.execute(select(MLSettings).where(MLSettings.id == 1)).scalar_one()
    w = [0.0] * 1152
    w[slot] = 1.0
    db.add(TagHead(
        tag_id=tag_id, embedding_version=s.embedder_model_version,
        weights=w, bias=0.0, suggest_threshold=0.5, auto_apply_threshold=threshold,
        n_pos=n_pos, n_neg=90, ap=0.9, precision_cv=0.98, recall=0.7,
    ))


def _run(db, dry_run=False) -> HeadAutoApplyRun:
    run = HeadAutoApplyRun(dry_run=dry_run, params={"dry_run": dry_run}, status="running")
    db.add(run)
    db.flush()
    return run


def _applied_source(db, image_id, tag_id):
    return db.execute(
        select(image_tag.c.source)
        .where(image_tag.c.image_record_id == image_id)
        .where(image_tag.c.tag_id == tag_id)
    ).scalar_one_or_none()


def test_sweep_applies_to_matching_image(db_sync):
    img = _img(db_sync, "a" * 64, _emb(0))
    tag = Tag(name="autotag", kind=TagKind.general)
    db_sync.add(tag)
    db_sync.flush()
    _head(db_sync, tag.id, 0)
    run = _run(db_sync)
    db_sync.commit()
    result = auto_apply_sweep(db_sync, run, dry_run=False)
    assert result["n_applied"] == 1
    assert _applied_source(db_sync, img.id, tag.id) == "head_auto"


def test_sweep_dry_run_counts_but_writes_nothing(db_sync):
    img = _img(db_sync, "b" * 64, _emb(0))
    tag = Tag(name="previewtag", kind=TagKind.general)
    db_sync.add(tag)
    db_sync.flush()
    _head(db_sync, tag.id, 0)
    run = _run(db_sync, dry_run=True)
    db_sync.commit()
    result = auto_apply_sweep(db_sync, run, dry_run=True)
    assert result["n_applied"] == 1                      # it WOULD apply
    assert _applied_source(db_sync, img.id, tag.id) is None  # but wrote nothing


def test_sweep_skips_under_supported_head(db_sync):
    # n_pos below head_auto_apply_min_positives (default 30) → a precise-looking
    # but under-supported head never fires.
    img = _img(db_sync, "c" * 64, _emb(0))
    tag = Tag(name="weaktag", kind=TagKind.general)
    db_sync.add(tag)
    db_sync.flush()
    _head(db_sync, tag.id, 0, n_pos=5)
    run = _run(db_sync)
    db_sync.commit()
    result = auto_apply_sweep(db_sync, run, dry_run=False)
    assert result["n_applied"] == 0
    assert _applied_source(db_sync, img.id, tag.id) is None


def test_sweep_skips_ungraduated_head(db_sync):
    # auto_apply_threshold is None (head never reached the precision bar).
    _img(db_sync, "d" * 64, _emb(0))
    tag = Tag(name="nograd", kind=TagKind.general)
    db_sync.add(tag)
    db_sync.flush()
    _head(db_sync, tag.id, 0, threshold=None)
    run = _run(db_sync)
    db_sync.commit()
    result = auto_apply_sweep(db_sync, run, dry_run=False)
    assert result["n_applied"] == 0


@pytest.mark.asyncio
async def test_auto_apply_disabled_blocks_real_run(client, db):
    # head_auto_apply_enabled defaults False → a real sweep is refused (400).
    resp = await client.post("/api/heads/auto-apply", json={"dry_run": False})
    assert resp.status_code == 400
    assert (await resp.get_json())["error"] == "auto_apply_disabled"


@pytest.mark.asyncio
async def test_auto_apply_dry_run_allowed_when_disabled(client, db, monkeypatch):
    monkeypatch.setattr(
        "backend.app.tasks.ml.apply_head_tags.delay", lambda *a, **k: None
    )
    resp = await client.post("/api/heads/auto-apply", json={"dry_run": True})
    assert resp.status_code == 202
    assert (await resp.get_json())["status"] == "running"


@pytest.mark.asyncio
async def test_auto_apply_conflict_when_one_running(client, db, monkeypatch):
    monkeypatch.setattr(
        "backend.app.tasks.ml.apply_head_tags.delay", lambda *a, **k: None
    )
    db.add(HeadAutoApplyRun(dry_run=True, params={}, status="running"))
    await db.flush()
    await db.commit()
    resp = await client.post("/api/heads/auto-apply", json={"dry_run": True})
    assert resp.status_code == 409
    assert (await resp.get_json())["error"] == "auto_apply_already_running"