FabledCurator/tests/test_external_worker.py

"""Integration tests for the external-link download worker (tasks/external)."""

import io
import zipfile

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

import backend.app.tasks.external as ext
from backend.app.models import (
    Artist,
    ExternalLink,
    ImageProvenance,
    Post,
    PostAttachment,
    Source,
)
from backend.app.services.external_fetch import FetchResult

pytestmark = pytest.mark.integration


def _structured_jpeg_bytes():
    # A non-uniform image so it survives the import filters (a solid color would
    # phash-collapse / fail thresholds).
    im = Image.new("RGB", (256, 256))
    px = im.load()
    for y in range(256):
        for x in range(256):
            px[x, y] = (x, y, (x + y) % 256)
    buf = io.BytesIO()
    im.save(buf, "JPEG")
    return buf.getvalue()


class _FakeLock:
    def acquire(self, blocking=False):
        return True

    def release(self):
        pass


class _FakeRedis:
    def lock(self, name, timeout=None, blocking=False):
        return _FakeLock()


def _seed(db_sync, *, host="pixeldrain", status="pending"):
    # Key per host so a single test can seed several (artist name is UNIQUE).
    artist = Artist(name=f"Ext {host}", slug=f"ext-{host}")
    db_sync.add(artist)
    db_sync.flush()
    source = Source(
        artist_id=artist.id, platform="patreon",
        url=f"https://patreon.com/{host}", enabled=True, config_overrides={},
    )
    db_sync.add(source)
    db_sync.flush()
    post = Post(
        source_id=source.id, artist_id=artist.id, external_post_id=f"EXT-{host}",
        post_url=f"https://patreon.com/posts/{host}",
    )
    db_sync.add(post)
    db_sync.flush()
    link = ExternalLink(
        post_id=post.id, artist_id=artist.id, host=host,
        url=f"https://{host}.test/file", status=status,
    )
    db_sync.add(link)
    db_sync.commit()
    return post, link


def test_fetch_external_link_downloads_and_attaches(db_sync, tmp_path, monkeypatch):
    post, link = _seed(db_sync)
    link_id = link.id

    monkeypatch.setattr(ext, "IMAGES_ROOT", tmp_path)
    monkeypatch.setattr(ext, "_redis", lambda: _FakeRedis())

    def fake_fetch(host, url, dest_dir, *, timeout, should_stop=lambda: False):
        dest_dir.mkdir(parents=True, exist_ok=True)
        f = dest_dir / "film.bin"  # non-art → PostAttachment (no thumb/ML enqueue)
        f.write_bytes(b"a film pack")
        return FetchResult(files=[f], bytes=f.stat().st_size)

    monkeypatch.setattr(ext, "fetch_external", fake_fetch)

    out = ext.fetch_external_link(link_id)
    assert out.get("files") == 1

    db_sync.expire_all()
    refreshed = db_sync.get(ExternalLink, link_id)
    assert refreshed.status == "downloaded"
    assert refreshed.completed_at is not None
    # The file was captured as a PostAttachment linked to the SAME post.
    atts = db_sync.execute(
        select(PostAttachment).where(PostAttachment.post_id == post.id)
    ).scalars().all()
    assert len(atts) == 1


def test_fetch_external_link_records_failure(db_sync, tmp_path, monkeypatch):
    _, link = _seed(db_sync)
    link_id = link.id
    monkeypatch.setattr(ext, "IMAGES_ROOT", tmp_path)
    monkeypatch.setattr(ext, "_redis", lambda: _FakeRedis())
    monkeypatch.setattr(
        ext, "fetch_external",
        lambda *a, **k: FetchResult(error="host 503"),
    )

    ext.fetch_external_link(link_id)
    db_sync.expire_all()
    refreshed = db_sync.get(ExternalLink, link_id)
    assert refreshed.status == "failed"
    assert refreshed.attempts == 1
    assert "503" in refreshed.last_error


def test_fetch_external_link_dead_letters_at_threshold(db_sync, tmp_path, monkeypatch):
    _, link = _seed(db_sync)
    link.attempts = ext.DEAD_LETTER_THRESHOLD - 1
    db_sync.commit()
    link_id = link.id
    monkeypatch.setattr(ext, "IMAGES_ROOT", tmp_path)
    monkeypatch.setattr(ext, "_redis", lambda: _FakeRedis())
    monkeypatch.setattr(ext, "fetch_external", lambda *a, **k: FetchResult(error="nope"))

    ext.fetch_external_link(link_id)
    db_sync.expire_all()
    assert db_sync.get(ExternalLink, link_id).status == "dead"


def test_fetch_external_link_skips_non_claimable(db_sync, tmp_path, monkeypatch):
    _, link = _seed(db_sync, status="downloaded")
    link_id = link.id
    called = []
    monkeypatch.setattr(ext, "_redis", lambda: _FakeRedis())
    monkeypatch.setattr(
        ext, "fetch_external",
        lambda *a, **k: called.append(1) or FetchResult(),
    )
    out = ext.fetch_external_link(link_id)
    assert out.get("skipped") == "not claimable"
    assert called == []  # never fetched


def test_sweep_enqueues_pending_and_retryable(db_sync, monkeypatch):
    _, l1 = _seed(db_sync, host="pixeldrain")
    _, l2 = _seed(db_sync, host="mega", status="failed")
    # A dead one must NOT be swept.
    _, l3 = _seed(db_sync, host="dropbox", status="dead")

    enqueued = []
    monkeypatch.setattr(ext.fetch_external_link, "delay", lambda lid: enqueued.append(lid))

    out = ext.sweep_external_links()
    assert out["enqueued"] == 2
    assert set(enqueued) == {l1.id, l2.id}


def test_downloaded_archive_gets_provenance_and_tagging(db_sync, tmp_path, monkeypatch):
    """An archive downloaded by the worker is extracted, provenance-linked to the
    post, AND queued for tagging + thumbnails — exactly like an extracted zip on
    the normal download path."""
    post, link = _seed(db_sync, host="mega")
    link_id = link.id

    monkeypatch.setattr(ext, "IMAGES_ROOT", tmp_path)
    monkeypatch.setattr(ext, "_redis", lambda: _FakeRedis())

    def fake_fetch(host, url, dest_dir, *, timeout, should_stop=lambda: False):
        dest_dir.mkdir(parents=True, exist_ok=True)
        zpath = dest_dir / "pack.zip"
        with zipfile.ZipFile(zpath, "w") as z:
            z.writestr("art.jpg", _structured_jpeg_bytes())
        return FetchResult(files=[zpath], bytes=zpath.stat().st_size)

    monkeypatch.setattr(ext, "fetch_external", fake_fetch)

    tagged, thumbed = [], []
    import backend.app.tasks.ml as ml_mod
    import backend.app.tasks.thumbnail as thumb_mod
    monkeypatch.setattr(ml_mod.tag_and_embed, "delay", lambda i: tagged.append(i))
    monkeypatch.setattr(thumb_mod.generate_thumbnail, "delay", lambda i: thumbed.append(i))

    out = ext.fetch_external_link(link_id)
    assert out["images"] >= 1

    db_sync.expire_all()
    assert db_sync.get(ExternalLink, link_id).status == "downloaded"
    # Extracted member is provenance-linked to the SAME post (as extracted zips).
    prov = db_sync.execute(
        select(ImageProvenance).where(ImageProvenance.post_id == post.id)
    ).scalars().all()
    assert len(prov) >= 1
    member_ids = {p.image_record_id for p in prov}
    # Tagging + thumbnails queued for exactly those member images.
    assert set(tagged) == member_ids
    assert set(thumbed) == member_ids