FabledCurator/tests/test_ml_allowlist.py

import pytest
from sqlalchemy import select

from backend.app.models import (
    ImagePrediction,
    TagAlias,
    TagAllowlist,
    TagKind,
    TagSuggestionRejection,
)
from backend.app.models.tag import image_tag
from backend.app.services.ml.allowlist import AllowlistService
from backend.app.services.tag_service import TagService

pytestmark = pytest.mark.integration


async def _make_image(db, sha: str = "x" * 64):
    from backend.app.models import ImageRecord
    img = ImageRecord(
        path=f"/images/{sha[:8]}.jpg", sha256=sha, size_bytes=1,
        mime="image/jpeg", width=1, height=1,
        origin="imported_filesystem", integrity_status="unknown",
    )
    db.add(img)
    await db.flush()
    return img


async def _add_pred(db, image_id, raw_name, score, category="general"):
    db.add(ImagePrediction(
        image_record_id=image_id, raw_name=raw_name,
        category=category, score=score,
    ))
    await db.flush()


@pytest.mark.asyncio
async def test_accept_applies_and_allowlists(db):
    img = await _make_image(db)
    tag = await TagService(db).find_or_create("Hero", TagKind.character)
    svc = AllowlistService(db)
    newly_added = await svc.accept(img.id, tag.id)
    assert newly_added is True

    applied = (
        await db.execute(
            select(image_tag.c.source)
            .where(image_tag.c.image_record_id == img.id)
            .where(image_tag.c.tag_id == tag.id)
        )
    ).scalar_one()
    assert applied == "ml_accepted"
    assert await db.get(TagAllowlist, tag.id) is not None


@pytest.mark.asyncio
async def test_accept_idempotent_allowlist(db):
    img = await _make_image(db)
    tag = await TagService(db).find_or_create("Hero2", TagKind.character)
    svc = AllowlistService(db)
    assert await svc.accept(img.id, tag.id) is True
    assert await svc.accept(img.id, tag.id) is False


@pytest.mark.asyncio
async def test_reject_applied_tag_records_rejection(db):
    img = await _make_image(db)
    tag = await TagService(db).find_or_create("Removeme", TagKind.general)
    svc = AllowlistService(db)
    await svc.accept(img.id, tag.id)
    await svc.reject_applied_tag(img.id, tag.id)

    still_applied = (
        await db.execute(
            select(image_tag.c.tag_id)
            .where(image_tag.c.image_record_id == img.id)
            .where(image_tag.c.tag_id == tag.id)
        )
    ).scalar_one_or_none()
    assert still_applied is None
    rej = await db.get(TagSuggestionRejection, (img.id, tag.id))
    assert rej is not None


@pytest.mark.asyncio
async def test_dismiss_records_rejection(db):
    img = await _make_image(db)
    tag = await TagService(db).find_or_create("Dismissme", TagKind.general)
    await AllowlistService(db).dismiss(img.id, tag.id)
    assert await db.get(TagSuggestionRejection, (img.id, tag.id)) is not None


@pytest.mark.asyncio
async def test_add_alias_and_accept(db):
    img = await _make_image(db)
    canonical = await TagService(db).find_or_create(
        "Canonical Char", TagKind.character
    )
    svc = AllowlistService(db)
    await svc.add_alias_and_accept(
        img.id, "model_char_name", "character", canonical.id
    )
    from backend.app.services.ml.aliases import AliasService
    resolved = await AliasService(db).resolve("model_char_name", "character")
    assert resolved.id == canonical.id
    assert await db.get(TagAllowlist, canonical.id) is not None


@pytest.mark.asyncio
async def test_update_threshold_and_remove(db):
    tag = await TagService(db).find_or_create("Thr", TagKind.general)
    svc = AllowlistService(db)
    img = await _make_image(db)
    await svc.accept(img.id, tag.id)
    await svc.update_threshold(tag.id, 0.80)
    row = await db.get(TagAllowlist, tag.id)
    assert abs(row.min_confidence - 0.80) < 1e-6
    await svc.remove(tag.id)
    assert await db.get(TagAllowlist, tag.id) is None


@pytest.mark.asyncio
async def test_coverage_by_threshold_direct_name(db):
    tag = await TagService(db).find_or_create("Cov", TagKind.general)
    svc = AllowlistService(db)
    for i, score in enumerate((0.95, 0.80, 0.60)):
        img = await _make_image(db, sha=f"c{i:063d}")
        await _add_pred(db, img.id, "Cov", score)
    assert await svc.coverage(tag.id, 0.90) == 1
    assert await svc.coverage(tag.id, 0.70) == 2
    assert await svc.coverage(tag.id, 0.50) == 3


@pytest.mark.asyncio
async def test_coverage_via_alias_respects_category(db):
    tag = await TagService(db).find_or_create("Aliased", TagKind.character)
    db.add(TagAlias(
        alias_string="model_key", alias_category="character",
        canonical_tag_id=tag.id,
    ))
    await db.flush()
    svc = AllowlistService(db)
    hit = await _make_image(db, sha=f"a{0:063d}")
    await _add_pred(db, hit.id, "model_key", 0.92, category="character")
    # Same alias string but wrong category must NOT resolve to the tag.
    miss = await _make_image(db, sha=f"a{1:063d}")
    await _add_pred(db, miss.id, "model_key", 0.99, category="general")
    assert await svc.coverage(tag.id, 0.90) == 1


@pytest.mark.asyncio
async def test_list_all_reports_applied_and_coverage(db):
    tag = await TagService(db).find_or_create("Both", TagKind.general)
    svc = AllowlistService(db)
    applied_img = await _make_image(db, sha=f"b{0:063d}")
    await svc.accept(applied_img.id, tag.id)  # applies + allowlists
    await _add_pred(db, applied_img.id, "Both", 0.95)
    # A second image only has a qualifying prediction (covered, not applied).
    cov_img = await _make_image(db, sha=f"b{1:063d}")
    await _add_pred(db, cov_img.id, "Both", 0.95)

    rows = await svc.list_all()
    row = next(r for r in rows if r.tag_id == tag.id)
    assert row.applied_count == 1                 # only the accepted image
    assert row.coverage_count == 2                # both have a ≥threshold pred


@pytest.mark.asyncio
async def test_update_threshold_clamped_to_store_floor(db):
    # A min_confidence below the store floor (default 0.70) is clamped up —
    # predictions below the floor aren't stored, so a lower threshold can't
    # apply more permissively than the floor (#764).
    tag = await TagService(db).find_or_create("Lowthr", TagKind.general)
    svc = AllowlistService(db)
    img = await _make_image(db)
    await svc.accept(img.id, tag.id)
    await svc.update_threshold(tag.id, 0.30)
    row = await db.get(TagAllowlist, tag.id)
    assert abs(row.min_confidence - 0.70) < 1e-6