109 lines
3.3 KiB
Python
109 lines
3.3 KiB
Python
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", [])]
|
|
assert "sword" in names
|
|
assert "lowconf" not in names
|
|
|
|
|
|
@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"]
|
|
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"]
|