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
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 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, color in enumerate(("red", "blue")):
        arc = store_dir / f"{n}_archive_{color}"
        with zipfile.ZipFile(arc, "w") as zf:
            zf.writestr(f"{color}.jpg", _jpeg(color))
        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