FabledCurator/tests/test_suggestions_bulk.py

"""Consensus (for_selection) over the tagging-v2 HEAD suggestion source."""
import pytest
from sqlalchemy import select

from backend.app import create_app
from backend.app.models import ImageRecord, MLSettings, TagHead, TagKind
from backend.app.models.tag import image_tag
from backend.app.services.ml.suggestions import SuggestionService
from backend.app.services.tag_service import TagService

pytestmark = pytest.mark.integration


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


async def _img(db, sha: str, emb=None) -> 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)
    await db.flush()
    return img


async def _head(db, tag_id: int, slot: int = 0):
    s = (await db.execute(select(MLSettings).where(MLSettings.id == 1))).scalar_one()
    weights = [0.0] * 1152
    weights[slot] = 1.0
    db.add(TagHead(
        tag_id=tag_id, embedding_version=s.embedder_model_version,
        weights=weights, bias=0.0, suggest_threshold=0.5,
        auto_apply_threshold=None, n_pos=10, n_neg=30,
        ap=0.8, precision_cv=0.9, recall=0.6,
    ))


@pytest.mark.asyncio
async def test_consensus_includes_tag_over_threshold(db):
    t = await TagService(db).find_or_create("sword", TagKind.general)
    a = await _img(db, "a" * 64, _emb(0))
    b = await _img(db, "b" * 64, _emb(0))
    await _head(db, t.id, slot=0)
    await db.commit()
    res = await SuggestionService(db).for_selection([a.id, b.id], threshold=0.8)
    s = next(s for s in res["general"] if s["canonical_tag_id"] == t.id)
    assert s["coverage"] == 1.0          # suggested on both
    assert s["confidence"] > 0.5


@pytest.mark.asyncio
async def test_consensus_counts_already_applied_for_coverage(db):
    t = await TagService(db).find_or_create("sky", TagKind.general)
    a = await _img(db, "c" * 64, _emb(0))   # head suggests it
    b = await _img(db, "d" * 64, None)       # no embedding; tag applied instead
    await _head(db, t.id, slot=0)
    await db.execute(
        image_tag.insert().values(
            image_record_id=b.id, tag_id=t.id, source="manual"
        )
    )
    await db.commit()
    res = await SuggestionService(db).for_selection([a.id, b.id], threshold=0.8)
    s = next(s for s in res["general"] if s["canonical_tag_id"] == t.id)
    assert s["coverage"] == 1.0          # 1 suggested + 1 applied / 2


@pytest.mark.asyncio
async def test_consensus_excludes_below_threshold(db):
    t = await TagService(db).find_or_create("rare", TagKind.general)
    a = await _img(db, "e" * 64, _emb(0))   # suggested here
    b = await _img(db, "f" * 64, None)       # not here → coverage 0.5 < 0.8
    await _head(db, t.id, slot=0)
    await db.commit()
    res = await SuggestionService(db).for_selection([a.id, b.id], threshold=0.8)
    assert all(s["name"] != "rare" for s in res.get("general", []))


@pytest.mark.asyncio
async def test_consensus_threshold_clamped_and_empty_for_no_ids(db):
    res = await SuggestionService(db).for_selection([], threshold=5.0)
    assert res == {}


@pytest.mark.asyncio
async def test_bulk_suggestions_route(db):
    t = await TagService(db).find_or_create("sword", TagKind.general)
    a = await _img(db, "i" * 64, _emb(0))
    await _head(db, t.id, slot=0)
    await db.commit()
    app = create_app()
    async with app.test_client() as c:
        resp = await c.post(
            "/api/suggestions/bulk", json={"image_ids": [a.id]}
        )
        assert resp.status_code == 200
        body = await resp.get_json()
        assert body["evaluated"] == 1
        assert body["threshold"] == 0.8
        assert "suggestions" in body


@pytest.mark.asyncio
async def test_bulk_suggestions_requires_ids(db):
    app = create_app()
    async with app.test_client() as c:
        resp = await c.post("/api/suggestions/bulk", json={})
        assert resp.status_code == 400