FabledCurator/tests/test_tasks_ml.py

"""tag_and_embed / backfill task tests. Models aren't in CI, so we test
the pure helpers (_aggregate_video_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 _aggregate_video_predictions, _is_video


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 _pred(name, conf, cat="general"):
    return {name: TagPrediction(name, cat, conf)}


def test_aggregate_video_keeps_corroborated_and_means():
    # #747: 4 frames; "smile" in 3, "sword" in 1 (noise). min_frames=2.
    per_frame = [
        {"smile": TagPrediction("smile", "general", 0.6),
         "sword": TagPrediction("sword", "general", 0.9)},
        _pred("smile", 0.8),
        _pred("smile", 0.7),
        {},
    ]
    out = _aggregate_video_predictions(per_frame, min_frames=2)
    assert "sword" not in out  # one-frame flicker dropped
    assert abs(out["smile"]["confidence"] - (0.6 + 0.8 + 0.7) / 3) < 1e-9  # mean, not max


def test_aggregate_video_clamps_min_frames_to_sample_count():
    # Short video: 1 frame but min_frames=3 — clamp so it still tags.
    out = _aggregate_video_predictions([_pred("solo", 0.8)], min_frames=3)
    assert out["solo"]["confidence"] == 0.8


def test_aggregate_video_empty():
    assert _aggregate_video_predictions([], min_frames=3) == {}


@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