127 lines
4.0 KiB
Python
127 lines
4.0 KiB
Python
"""tag_and_embed / backfill task tests. Models aren't in CI, so we test
|
|
the pure helpers (_maxpool_predictions, _is_video) as unit tests, and the
|
|
DB-touching backfill query as an integration test with monkeypatched
|
|
inference.
|
|
"""
|
|
|
|
from pathlib import Path
|
|
|
|
import pytest
|
|
|
|
from backend.app.services.ml.tagger import TagPrediction
|
|
from backend.app.tasks.ml import _is_video, _maxpool_predictions
|
|
|
|
|
|
def test_is_video():
|
|
assert _is_video(Path("a.mp4")) is True
|
|
assert _is_video(Path("a.MKV")) is True
|
|
assert _is_video(Path("a.jpg")) is False
|
|
|
|
|
|
def test_maxpool_predictions():
|
|
f1 = {"smile": TagPrediction("smile", "general", 0.6)}
|
|
f2 = {
|
|
"smile": TagPrediction("smile", "general", 0.9),
|
|
"sword": TagPrediction("sword", "general", 0.7),
|
|
}
|
|
merged = _maxpool_predictions([f1, f2])
|
|
assert merged["smile"]["confidence"] == 0.9
|
|
assert merged["sword"]["confidence"] == 0.7
|
|
|
|
|
|
@pytest.mark.integration
|
|
@pytest.mark.asyncio
|
|
async def test_backfill_enqueues_missing(db, monkeypatch):
|
|
from backend.app.models import ImageRecord
|
|
from backend.app.tasks import ml as ml_tasks
|
|
|
|
calls = []
|
|
monkeypatch.setattr(
|
|
ml_tasks.tag_and_embed, "delay", lambda image_id: calls.append(image_id)
|
|
)
|
|
|
|
img = ImageRecord(
|
|
path="/images/n.jpg", sha256="n" * 64, size_bytes=1,
|
|
mime="image/jpeg", width=1, height=1,
|
|
origin="imported_filesystem", integrity_status="unknown",
|
|
tagger_predictions=None, siglip_embedding=None,
|
|
)
|
|
db.add(img)
|
|
await db.commit()
|
|
|
|
count = ml_tasks.backfill()
|
|
assert count >= 1
|
|
assert img.id in calls
|
|
|
|
|
|
@pytest.mark.integration
|
|
@pytest.mark.asyncio
|
|
async def test_apply_allowlist_applies_above_threshold(db):
|
|
from sqlalchemy import select
|
|
|
|
from backend.app.models import ImageRecord, TagAllowlist, TagKind
|
|
from backend.app.models.tag import image_tag
|
|
from backend.app.services.tag_service import TagService
|
|
from backend.app.tasks import ml as ml_tasks
|
|
from tests._prediction_helpers import seed_predictions
|
|
|
|
tag = await TagService(db).find_or_create("autohero", TagKind.character)
|
|
db.add(TagAllowlist(tag_id=tag.id, min_confidence=0.95))
|
|
img = ImageRecord(
|
|
path="/images/al.jpg", sha256="al" + "0" * 62, 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, {"autohero": {"category": "character", "confidence": 0.97}}
|
|
)
|
|
await db.commit()
|
|
|
|
n = ml_tasks.apply_allowlist_tags(tag_id=tag.id)
|
|
assert n >= 1
|
|
src = (
|
|
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 src == "ml_auto"
|
|
|
|
|
|
@pytest.mark.integration
|
|
@pytest.mark.asyncio
|
|
async def test_apply_allowlist_skips_below_threshold(db):
|
|
from sqlalchemy import select
|
|
|
|
from backend.app.models import ImageRecord, TagAllowlist, TagKind
|
|
from backend.app.models.tag import image_tag
|
|
from backend.app.services.tag_service import TagService
|
|
from backend.app.tasks import ml as ml_tasks
|
|
from tests._prediction_helpers import seed_predictions
|
|
|
|
tag = await TagService(db).find_or_create("lowconf", TagKind.character)
|
|
db.add(TagAllowlist(tag_id=tag.id, min_confidence=0.95))
|
|
img = ImageRecord(
|
|
path="/images/lc.jpg", sha256="lc" + "0" * 62, 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, {"lowconf": {"category": "character", "confidence": 0.40}}
|
|
)
|
|
await db.commit()
|
|
ml_tasks.apply_allowlist_tags(tag_id=tag.id)
|
|
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 applied is None
|