FabledCurator/tests/test_reextract_archives.py

"""#713 part 2: re-extract archive PostAttachments that were filed opaquely
(magic-byte gate missed them) and link their members to the post."""

import hashlib
import io
import zipfile

import pytest
from PIL import Image, ImageDraw
from sqlalchemy import select

from backend.app.models import (
    Artist,
    ImageProvenance,
    ImageRecord,
    Post,
    PostAttachment,
    Source,
)
from backend.app.services import cleanup_service

pytestmark = pytest.mark.integration


def _jpeg(color, size=256):
    buf = io.BytesIO()
    Image.new("RGB", (size, size), color).save(buf, "JPEG")
    return buf.getvalue()


def _patterned_jpeg(idx, size=256):
    """Structurally distinct per idx so the members don't phash-collapse — solid
    colors collide to distance 0 and the second would dedupe away (the two-axis
    dedup gotcha). A rectangle at an idx-derived position keeps phashes apart."""
    img = Image.new("RGB", (size, size), "white")
    x = 20 + idx * 60
    ImageDraw.Draw(img).rectangle([x, x, x + 70, x + 70], fill="black")
    buf = io.BytesIO()
    img.save(buf, "JPEG")
    return buf.getvalue()


def test_reextract_links_archive_members_to_post(db_sync, tmp_path, monkeypatch):
    from backend.app.tasks import ml as ml_mod
    from backend.app.tasks import thumbnail as thumb_mod

    # No broker in this path — the post-import enqueue is best-effort anyway.
    monkeypatch.setattr(thumb_mod.generate_thumbnail, "delay", lambda *a, **k: None)
    monkeypatch.setattr(ml_mod.tag_and_embed, "delay", lambda *a, **k: None)

    images_root = tmp_path / "images"
    images_root.mkdir()

    artist = Artist(name="Bob", slug="bob")
    db_sync.add(artist)
    db_sync.flush()
    source = Source(
        artist_id=artist.id, platform="patreon",
        url="https://patreon.com/bob", enabled=True, config_overrides={},
    )
    db_sync.add(source)
    db_sync.flush()
    post = Post(
        source_id=source.id, artist_id=artist.id, external_post_id="59102952",
        post_url="https://www.patreon.com/posts/59102952",
    )
    db_sync.add(post)
    db_sync.flush()

    # A real zip stored under a mangled / extension-less name (the failure case).
    store_dir = images_root / "attachments" / "abc"
    store_dir.mkdir(parents=True)
    arc = store_dir / "01_https___www.patreon.com_media-u_v3_59102952"
    with zipfile.ZipFile(arc, "w") as zf:
        zf.writestr("inside.jpg", _jpeg("red"))
    sha = hashlib.sha256(arc.read_bytes()).hexdigest()
    db_sync.add(PostAttachment(
        post_id=post.id, artist_id=artist.id, sha256=sha, path=str(arc),
        original_filename=arc.name, ext="", size_bytes=arc.stat().st_size,
    ))
    db_sync.commit()

    summary = cleanup_service.reextract_archive_attachments(
        db_sync, images_root=images_root,
    )
    assert summary["archives"] == 1
    assert summary["members_imported"] == 1
    assert summary["posts_touched"] == 1

    images = db_sync.execute(select(ImageRecord)).scalars().all()
    assert len(images) == 1
    prov = db_sync.execute(
        select(ImageProvenance).where(ImageProvenance.post_id == post.id)
    ).scalars().all()
    assert len(prov) == 1
    assert prov[0].image_record_id == images[0].id

    # Idempotent — a second run imports nothing new (member dedups by sha256).
    again = cleanup_service.reextract_archive_attachments(
        db_sync, images_root=images_root,
    )
    assert again["members_imported"] == 0
    assert db_sync.execute(select(ImageRecord)).scalars().all() == images


def test_reextract_timebox_resumes_from_cursor(db_sync, tmp_path, monkeypatch):
    """A 0-second budget cuts the chunk after the first attachment and reports a
    resume cursor; the next run starts strictly after it and finishes the rest."""
    from backend.app.tasks import ml as ml_mod
    from backend.app.tasks import thumbnail as thumb_mod

    monkeypatch.setattr(thumb_mod.generate_thumbnail, "delay", lambda *a, **k: None)
    monkeypatch.setattr(ml_mod.tag_and_embed, "delay", lambda *a, **k: None)

    images_root = tmp_path / "images"
    images_root.mkdir()

    artist = Artist(name="Bob", slug="bob")
    db_sync.add(artist)
    db_sync.flush()
    source = Source(
        artist_id=artist.id, platform="patreon",
        url="https://patreon.com/bob", enabled=True, config_overrides={},
    )
    db_sync.add(source)
    db_sync.flush()
    post = Post(
        source_id=source.id, artist_id=artist.id, external_post_id="42",
        post_url="https://www.patreon.com/posts/42",
    )
    db_sync.add(post)
    db_sync.flush()

    store_dir = images_root / "attachments" / "two"
    store_dir.mkdir(parents=True)
    att_ids = []
    for n in range(2):
        arc = store_dir / f"{n}_archive"
        with zipfile.ZipFile(arc, "w") as zf:
            zf.writestr(f"member{n}.jpg", _patterned_jpeg(n))
        sha = hashlib.sha256(arc.read_bytes()).hexdigest()
        att = PostAttachment(
            post_id=post.id, artist_id=artist.id, sha256=sha, path=str(arc),
            original_filename=arc.name, ext="", size_bytes=arc.stat().st_size,
        )
        db_sync.add(att)
        db_sync.flush()
        att_ids.append(att.id)
    db_sync.commit()

    # Budget 0 → break right after the first attachment commits.
    first = cleanup_service.reextract_archive_attachments(
        db_sync, images_root=images_root, time_budget_seconds=0.0,
    )
    assert first["partial"] is True
    assert first["scanned"] == 1
    assert first["members_imported"] == 1
    assert first["resume_after_id"] == att_ids[0]

    # Resume strictly after the cursor — picks up the second, then runs dry.
    second = cleanup_service.reextract_archive_attachments(
        db_sync, images_root=images_root, after_id=att_ids[0],
    )
    assert second["partial"] is False
    assert second["scanned"] == 1
    assert second["members_imported"] == 1
    assert len(db_sync.execute(select(ImageRecord)).scalars().all()) == 2