FabledCurator/tests/test_api_suggestions.py

import pytest
from sqlalchemy import select

from backend.app.celery_app import celery
from backend.app.models import ImageRecord, MLSettings, TagHead, TagKind
from backend.app.services.tag_service import TagService

pytestmark = pytest.mark.integration


@pytest.fixture(autouse=True)
def eager():
    celery.conf.task_always_eager = True
    yield
    celery.conf.task_always_eager = False


async def _img(db, preds, sha="s" * 64):
    from tests._prediction_helpers import seed_predictions

    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",
    )
    db.add(img)
    await db.commit()
    await seed_predictions(db, img.id, preds)
    await db.commit()
    return img


@pytest.mark.asyncio
async def test_get_suggestions(client, db):
    # Suggestions come from a trained head now (Camie/centroid removed): an image
    # whose embedding aligns with the head surfaces that concept.
    s = (await db.execute(select(MLSettings).where(MLSettings.id == 1))).scalar_one()
    img = ImageRecord(
        path="/images/headsug.jpg", sha256="h" * 64, size_bytes=1,
        mime="image/jpeg", width=1, height=1, origin="imported_filesystem",
        integrity_status="unknown", siglip_embedding=[3.0] + [0.0] * 1151,
    )
    db.add(img)
    await db.flush()
    tag = await TagService(db).find_or_create("sword", TagKind.general)
    db.add(TagHead(
        tag_id=tag.id, embedding_version=s.embedder_model_version,
        weights=[1.0] + [0.0] * 1151, 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,
    ))
    await db.commit()
    resp = await client.get(f"/api/images/{img.id}/suggestions")
    assert resp.status_code == 200
    body = await resp.get_json()
    general = body["by_category"].get("general", [])
    s2 = next(x for x in general if x["canonical_tag_id"] == tag.id)
    assert s2["source"] == "head"


@pytest.mark.asyncio
async def test_accept_requires_tag_id(client, db):
    img = await _img(db, {})
    resp = await client.post(
        f"/api/images/{img.id}/suggestions/accept", json={}
    )
    assert resp.status_code == 400


@pytest.mark.asyncio
async def test_accept_then_applied(client, db):
    img = await _img(db, {})
    tag = await TagService(db).find_or_create("AcceptMe", TagKind.character)
    await db.commit()
    resp = await client.post(
        f"/api/images/{img.id}/suggestions/accept", json={"tag_id": tag.id}
    )
    assert resp.status_code == 200
    body = await resp.get_json()
    # #7b: a fresh accept newly-allowlists → projection payload for the toast.
    assert body["allowlisted"] is True
    assert body["tag_id"] == tag.id
    assert body["tag_name"] == "AcceptMe"
    assert "projected_count" in body


@pytest.mark.asyncio
async def test_accept_already_allowlisted_reports_not_new(client, db):
    img1 = await _img(db, {}, sha="c" * 64)
    img2 = await _img(db, {}, sha="d" * 64)
    tag = await TagService(db).find_or_create("Twice", TagKind.character)
    await db.commit()
    first = await client.post(
        f"/api/images/{img1.id}/suggestions/accept", json={"tag_id": tag.id}
    )
    assert (await first.get_json())["allowlisted"] is True
    second = await client.post(
        f"/api/images/{img2.id}/suggestions/accept", json={"tag_id": tag.id}
    )
    body = await second.get_json()
    assert body["allowlisted"] is False   # already on the allowlist
    assert "projected_count" not in body


@pytest.mark.asyncio
async def test_dismiss(client, db):
    img = await _img(db, {})
    tag = await TagService(db).find_or_create("DismissMe", TagKind.general)
    await db.commit()
    resp = await client.post(
        f"/api/images/{img.id}/suggestions/dismiss", json={"tag_id": tag.id}
    )
    assert resp.status_code == 204


@pytest.mark.asyncio
async def test_undismiss_reverses_rejection(client, db):
    img = await _img(db, {})
    tag = await TagService(db).find_or_create("UndismissMe", TagKind.general)
    await db.commit()
    await client.post(
        f"/api/images/{img.id}/suggestions/dismiss", json={"tag_id": tag.id}
    )
    resp = await client.post(
        f"/api/images/{img.id}/suggestions/undismiss", json={"tag_id": tag.id}
    )
    assert resp.status_code == 204
    # Idempotent: un-rejecting again (nothing to clear) is still a 204.
    resp2 = await client.post(
        f"/api/images/{img.id}/suggestions/undismiss", json={"tag_id": tag.id}
    )
    assert resp2.status_code == 204


@pytest.mark.asyncio
async def test_alias_requires_fields(client, db):
    img = await _img(db, {})
    resp = await client.post(
        f"/api/images/{img.id}/suggestions/alias", json={"alias_string": "x"}
    )
    assert resp.status_code == 400