Files
FabledCurator/tests/test_tasks_ml.py
T
bvandeusen 3610ba495f
CI / lint (push) Successful in 3s
CI / frontend-build (push) Successful in 23s
CI / backend-lint-and-test (push) Successful in 33s
CI / integration (push) Successful in 3m14s
feat(ml): drop image_record.tagger_predictions — image_prediction is sole store (#768 step 3)
Read cutover verified in prod (suggestions + allowlist read image_prediction;
backfill complete at 908k rows / 51k images). Removes the old JSON column and
everything that fed it:

- ImageRecord.tagger_predictions column removed; migration 0046 DROPs it.
  tagger_model_version kept as the "tagged / current?" signal the backfill
  sweep reads (needs-tagging check switched to tagger_model_version IS NULL).
- tag_and_embed no longer dual-writes the JSON — image_prediction is the only
  write path.
- importer re-import reset drops the JSON line (image_prediction rows are
  already deleted on re-import).
- Retired the one-time #768 backfill task + the #764 prune task, their admin
  endpoints, and their Maintenance cards (Backfill/PrunePredictionsCard).
- Tests seed/assert via image_prediction; stale column refs removed.

Disk reclaim is NOT automatic: DROP COLUMN is a catalog change. Run
`VACUUM FULL image_record` off-hours afterward to return the ~100 GB to the OS
so DB backups go small (#739). image_prediction (~90 MB) stays in pg_dump — it's
the source of truth now.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
2026-06-11 18:52:33 -04:00

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",
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