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.suggestions import SuggestionService
from backend.app.services.tag_service import TagService

pytestmark = pytest.mark.integration


def _img(sha: str, predictions: dict) -> 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",
        tagger_predictions=predictions,
    )


@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 = _img(
        "a" * 64,
        {
            "lowconf": {"category": "general", "confidence": 0.30},
            "sword": {"category": "general", "confidence": 0.97},
        },
    )
    db.add(img)
    await db.flush()
    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 = _img(
        "d" * 64,
        {
            "lowconf": {"category": "general", "confidence": 0.30},
            "sword": {"category": "general", "confidence": 0.97},
        },
    )
    db.add(img)
    await db.flush()
    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 = _img("e" * 64, {"safe": {"category": "rating", "confidence": 0.99}})
    db.add(img2)
    await db.flush()
    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 = _img(
        "b" * 64,
        {"safe": {"category": "rating", "confidence": 0.99}},
    )
    db.add(img)
    await db.flush()
    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 = _img(
        "c" * 64,
        {"uchiha_sasuke": {"category": "character", "confidence": 0.96}},
    )
    db.add(img)
    await db.flush()
    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


@pytest.mark.asyncio
async def test_raw_tag_creates_new(db):
    img = _img(
        "d" * 64,
        {"brand_new_tag": {"category": "character", "confidence": 0.96}},
    )
    db.add(img)
    await db.flush()
    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
    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 = _img(
        "e" * 64,
        {"alreadyhere": {"category": "character", "confidence": 0.96}},
    )
    db.add(img)
    await db.flush()
    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"]