Files
FabledCurator/tests/test_phash_dedup.py
T
2026-05-17 22:25:48 -04:00

175 lines
5.7 KiB
Python

"""pHash near-dup + supersession importer tests.
Mirrors tests/test_importer.py exactly: the Importer is sync, driven via
the db_sync savepoint Session; tests are plain `def` (not asyncio).
"""
from pathlib import Path
import pytest
from PIL import Image
from sqlalchemy import func, select
from backend.app.models import ImageRecord, ImportSettings, Tag, TagKind
from backend.app.models.tag import image_tag
from backend.app.services.importer import Importer, SkipReason
from backend.app.services.thumbnailer import Thumbnailer
pytestmark = pytest.mark.integration
@pytest.fixture
def import_layout(tmp_path):
import_root = tmp_path / "import"
images_root = tmp_path / "images"
import_root.mkdir()
images_root.mkdir()
return import_root, images_root
@pytest.fixture
def importer(db_sync, import_layout):
import_root, images_root = import_layout
settings = db_sync.execute(
select(ImportSettings).where(ImportSettings.id == 1)
).scalar_one()
return Importer(
session=db_sync,
images_root=images_root,
import_root=import_root,
thumbnailer=Thumbnailer(images_root=images_root),
settings=settings,
)
def _set_threshold(importer, n):
importer.settings.phash_threshold = n
importer.session.flush()
def _write(path: Path, color, size):
path.parent.mkdir(parents=True, exist_ok=True)
Image.new("RGB", size, color).save(path)
def _write_split(path: Path, orient, size):
"""Structured (half black / half white) image — solid colors all
phash-collapse to the same value, so a real threshold test needs DCT
content. Vertical vs horizontal split = structurally far apart."""
path.parent.mkdir(parents=True, exist_ok=True)
w, h = size
im = Image.new("L", size, 0)
px = im.load()
for y in range(h):
for x in range(w):
if (x / w if orient == "v" else y / h) >= 0.5:
px[x, y] = 255
im.convert("RGB").save(path)
def test_non_similar_imports_with_phash(importer, import_layout):
import_root, _ = import_layout
src = import_root / "a.png"
_write(src, (12, 200, 60), (300, 300))
r = importer.import_one(src)
assert r.status == "imported"
row = importer.session.get(ImageRecord, r.image_id)
assert row.phash is not None and len(row.phash) == 16
def test_larger_existing_skips_new_phash_dup(importer, import_layout):
import_root, _ = import_layout
big = import_root / "big.png"
_write(big, (40, 90, 160), (800, 800))
r1 = importer.import_one(big)
assert r1.status == "imported"
small = import_root / "small.png"
_write(small, (40, 90, 160), (200, 200)) # same look, smaller
r2 = importer.import_one(small)
assert r2.status == "skipped"
assert r2.skip_reason == SkipReason.duplicate_phash
assert r2.image_id == r1.image_id
count = importer.session.execute(
select(func.count()).select_from(ImageRecord)
).scalar_one()
assert count == 1
def test_smaller_existing_is_superseded(importer, import_layout):
import_root, _ = import_layout
small = import_root / "small.png"
_write(small, (200, 30, 30), (200, 200))
r1 = importer.import_one(small)
eid = r1.image_id
old = importer.session.get(ImageRecord, eid)
old_path = old.path
tag = Tag(name="keepme", kind=TagKind.general)
importer.session.add(tag)
importer.session.flush()
importer.session.execute(
image_tag.insert().values(
image_record_id=eid, tag_id=tag.id, source="manual"
)
)
old.tagger_predictions = {"x": 1}
old.siglip_embedding = [0.0] * 1152
old.integrity_status = "ok"
importer.session.commit()
big = import_root / "big.png"
_write(big, (200, 30, 30), (900, 900)) # same look, larger
r2 = importer.import_one(big)
assert r2.status == "superseded"
assert r2.image_id == eid
importer.session.expire_all()
row = importer.session.get(ImageRecord, eid)
assert row.id == eid
assert row.width == 900 and row.height == 900
assert row.path != old_path
assert row.phash is not None
assert row.integrity_status == "unknown"
assert row.tagger_predictions is None
assert row.siglip_embedding is None
linked = importer.session.execute(
select(image_tag.c.tag_id).where(
image_tag.c.image_record_id == eid
)
).scalars().all()
assert tag.id in linked
assert not Path(old_path).exists()
assert Path(row.path).exists()
def test_threshold_controls_match(importer, import_layout):
# Structurally distinct images (orthogonal splits) are far apart in
# phash space, so a tight threshold keeps them independent rather than
# collapsing the larger one into a supersede.
import_root, _ = import_layout
_set_threshold(importer, 0)
a = import_root / "a.png"
_write_split(a, "v", (200, 200))
importer.import_one(a)
b = import_root / "b.png"
_write_split(b, "h", (900, 900)) # different structure, larger
r = importer.import_one(b)
assert r.status == "imported" # threshold 0 + far → independent import
def test_import_task_maps_superseded_to_complete_and_requeues():
from backend.app.services.importer import ImportResult
from backend.app.tasks.import_file import _map_result_to_status
assert _map_result_to_status(
ImportResult(status="superseded", image_id=5)
) == ("complete", True)
assert _map_result_to_status(
ImportResult(status="imported", image_id=5)
) == ("complete", True)
assert _map_result_to_status(
ImportResult(status="skipped", skip_reason=SkipReason.duplicate_phash)
) == ("skipped", False)
assert _map_result_to_status(
ImportResult(status="failed", error="boom")
) == ("failed", False)