FabledCurator/tests/test_ml_suggestions.py

import pytest

from backend.app.models import ImageRecord, TagKind
from backend.app.models.tag import image_tag
from backend.app.services.ml.aliases import AliasService
from backend.app.services.ml.allowlist import AllowlistService
from backend.app.services.ml.suggestions import SuggestionService
from backend.app.services.tag_service import TagService

pytestmark = pytest.mark.integration


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


async def _seed_img(db, sha: str, predictions: dict) -> ImageRecord:
    """#768: create an image + seed its predictions into image_prediction
    (the read path's source), returning the flushed record."""
    from tests._prediction_helpers import seed_predictions

    img = _img(sha)
    db.add(img)
    await db.flush()
    await seed_predictions(db, img.id, predictions)
    return img


@pytest.mark.asyncio
async def test_threshold_filters_low_confidence_general(db):
    # Default general threshold is 0.50 (alembic 0029 lowered it from
    # 0.95). Use 0.30/0.60 to keep the test asserting threshold behavior
    # rather than the exact cutoff number.
    img = await _seed_img(
        db,
        "a" * 64,
        {
            "lowconf": {"category": "general", "confidence": 0.30},
            "sword": {"category": "general", "confidence": 0.97},
        },
    )
    sl = await SuggestionService(db).for_image(img.id)
    names = [s.display_name for s in sl.by_category.get("general", [])]
    # display_name is normalized (tag_name.normalize) before surfacing.
    assert "Sword" in names
    assert "Lowconf" not in names


@pytest.mark.asyncio
async def test_threshold_override_surfaces_low_confidence(db):
    # The typed-dropdown "show everything the model saw" mode: threshold_override
    # surfaces stored predictions below the configured threshold (in canonical
    # formatting) so they can be picked instead of hand-typed (2026-06-09).
    img = await _seed_img(
        db,
        "d" * 64,
        {
            "lowconf": {"category": "general", "confidence": 0.30},
            "sword": {"category": "general", "confidence": 0.97},
        },
    )
    sl = await SuggestionService(db).for_image(img.id, threshold_override=0.0)
    names = [s.display_name for s in sl.by_category.get("general", [])]
    assert "Sword" in names
    assert "Lowconf" in names  # below the configured threshold, surfaced anyway

    # Unsurfaced categories are still excluded even with the override.
    img2 = await _seed_img(
        db, "e" * 64, {"safe": {"category": "rating", "confidence": 0.99}}
    )
    sl2 = await SuggestionService(db).for_image(img2.id, threshold_override=0.0)
    assert "rating" not in sl2.by_category


@pytest.mark.asyncio
async def test_unsurfaced_category_dropped(db):
    img = await _seed_img(
        db,
        "b" * 64,
        {"safe": {"category": "rating", "confidence": 0.99}},
    )
    sl = await SuggestionService(db).for_image(img.id)
    assert "rating" not in sl.by_category


@pytest.mark.asyncio
async def test_alias_resolution(db):
    tags = TagService(db)
    canonical = await tags.find_or_create("Sasuke Uchiha", TagKind.character)
    await AliasService(db).create("uchiha_sasuke", "character", canonical.id)
    img = await _seed_img(
        db,
        "c" * 64,
        {"uchiha_sasuke": {"category": "character", "confidence": 0.96}},
    )
    sl = await SuggestionService(db).for_image(img.id)
    chars = sl.by_category["character"]
    assert len(chars) == 1
    assert chars[0].display_name == "Sasuke Uchiha"
    assert chars[0].canonical_tag_id == canonical.id
    assert chars[0].creates_new_tag is False
    # Surfaced via an alias on the raw model key — the UI marks it + offers undo.
    assert chars[0].via_alias is True
    assert chars[0].raw_name == "uchiha_sasuke"


@pytest.mark.asyncio
async def test_raw_tag_creates_new(db):
    img = await _seed_img(
        db,
        "d" * 64,
        {"brand_new_tag": {"category": "character", "confidence": 0.96}},
    )
    sl = await SuggestionService(db).for_image(img.id)
    chars = sl.by_category["character"]
    # display_name is the normalized Camie name (underscores -> spaces,
    # title-cased), not the raw vocab key.
    assert chars[0].display_name == "Brand New Tag"
    assert chars[0].creates_new_tag is True
    # Not aliased, but the raw key is carried so the modal can author one.
    assert chars[0].via_alias is False
    assert chars[0].raw_name == "brand_new_tag"
    assert chars[0].canonical_tag_id is None


@pytest.mark.asyncio
async def test_applied_tag_not_suggested(db):
    tags = TagService(db)
    tag = await tags.find_or_create("alreadyhere", TagKind.character)
    img = await _seed_img(
        db,
        "e" * 64,
        {"alreadyhere": {"category": "character", "confidence": 0.96}},
    )
    await db.execute(
        image_tag.insert().values(
            image_record_id=img.id, tag_id=tag.id, source="manual"
        )
    )
    sl = await SuggestionService(db).for_image(img.id)
    assert "character" not in sl.by_category or not sl.by_category["character"]


@pytest.mark.asyncio
async def test_rejected_tag_surfaced_flagged_then_reversible(db):
    # A dismissed suggestion is NOT dropped: it stays in the list flagged
    # rejected=True so the rail can show it + offer one-click un-reject
    # (visible, reversible rejection — operator-asked 2026-06-27). A live
    # suggestion sorts ahead of the rejected one.
    tags = TagService(db)
    rejected_tag = await tags.find_or_create("rejectme", TagKind.general)
    img = await _seed_img(
        db,
        "f" * 64,
        {
            "rejectme": {"category": "general", "confidence": 0.96},
            "keepme": {"category": "general", "confidence": 0.90},
        },
    )
    await AllowlistService(db).dismiss(img.id, rejected_tag.id)

    sl = await SuggestionService(db).for_image(img.id)
    general = sl.by_category["general"]
    by_name = {s.display_name: s for s in general}
    assert by_name["Rejectme"].rejected is True
    assert by_name["Keepme"].rejected is False
    # Live suggestions sort ahead of rejected ones regardless of score.
    assert general[-1].display_name == "Rejectme"

    # Un-reject reverts it to a live suggestion.
    await AllowlistService(db).undismiss(img.id, rejected_tag.id)
    sl2 = await SuggestionService(db).for_image(img.id)
    by_name2 = {s.display_name: s for s in sl2.by_category["general"]}
    assert by_name2["Rejectme"].rejected is False