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, update

import backend.app.tasks.external as ext
from backend.app.models import (
    Artist,
    ExternalLink,
    ImageProvenance,
    ImageRecord,
    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, **kwargs):
        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_refetch_same_link_keeps_canonical_file(db_sync, tmp_path, monkeypatch):
    """#859: re-fetching the SAME link (identical bytes into the same per-link
    staging dir) must NOT unlink the canonical file out from under its record.
    The first fetch imports the image in place; the second dedups on sha256, and
    the path-safe unlink recognises the file IS the record's backing file and
    keeps it — so the post never 404s on playback."""
    _, link = _seed(db_sync, host="pixeldrain")
    link_id = link.id

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

    def fake_fetch(host, url, dest_dir, **kwargs):
        dest_dir.mkdir(parents=True, exist_ok=True)
        f = dest_dir / "clip.jpg"  # art → ImageRecord (imported in place)
        f.write_bytes(_structured_jpeg_bytes())
        return FetchResult(files=[f], bytes=f.stat().st_size)

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

    import backend.app.tasks.ml as ml_mod
    import backend.app.tasks.thumbnail as thumb_mod
    monkeypatch.setattr(ml_mod.embed_image, "delay", lambda i: None)
    monkeypatch.setattr(thumb_mod.generate_thumbnail, "delay", lambda i: None)

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

    db_sync.expire_all()
    rec = db_sync.execute(select(ImageRecord)).scalars().one()
    canonical = tmp_path / rec.path  # absolute path collapses the join
    assert canonical.exists()

    # Re-arm the SAME link and fetch again — same staging dir, identical bytes.
    db_sync.execute(
        update(ExternalLink).where(ExternalLink.id == link_id)
        .values(status="pending", attempts=0, completed_at=None)
    )
    db_sync.commit()

    out2 = ext.fetch_external_link(link_id)
    assert out2["images"] == 0  # deduped, no new record

    db_sync.expire_all()
    # Exactly one record still, and its file survived the dedup unlink (the bug).
    assert db_sync.execute(select(func.count(ImageRecord.id))).scalar() == 1
    assert canonical.exists()


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, **kwargs):
        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.embed_image, "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


def _link_dir(tmp_path, host, link_id):
    return tmp_path / f"ext-{host}" / "patreon" / "external" / f"EXT-{host}" / str(link_id)


def test_extracted_archive_is_not_left_in_the_link_dir(db_sync, tmp_path, monkeypatch):
    """An archive whose images were extracted is in the attachments store, so
    the downloaded copy is removed instead of sitting beside it."""
    _, link = _seed(db_sync, host="mega")
    link_id = link.id
    monkeypatch.setattr(ext, "IMAGES_ROOT", tmp_path)
    monkeypatch.setattr(ext, "_redis", lambda: _FakeRedis())
    written = []

    def fake_fetch(host, url, dest_dir, **kwargs):
        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())
        written.append(zpath)
        return FetchResult(files=[zpath], bytes=zpath.stat().st_size)

    monkeypatch.setattr(ext, "fetch_external", fake_fetch)
    import backend.app.tasks.ml as ml_mod
    import backend.app.tasks.thumbnail as thumb_mod
    monkeypatch.setattr(ml_mod.embed_image, "delay", lambda i: None)
    monkeypatch.setattr(thumb_mod.generate_thumbnail, "delay", lambda i: None)

    out = ext.fetch_external_link(link_id)
    assert out["images"] >= 1
    assert not written[0].exists()
    db_sync.expire_all()
    assert db_sync.execute(
        select(func.count(PostAttachment.id))
        .where(PostAttachment.original_filename == "pack.zip")
    ).scalar_one() == 1


def test_a_failed_fetch_leaves_no_partial_files(db_sync, tmp_path, monkeypatch):
    """A fetch that fails midway (megatools leaves `.megatmp.*` behind) clears
    the link's staging dir, so a dead link leaves nothing on disk."""
    _, 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, **kwargs):
        dest_dir.mkdir(parents=True, exist_ok=True)
        (dest_dir / ".megatmp.AbCd").write_bytes(b"half a game build")
        return FetchResult(error="megatools dl exit 1: timed out")

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

    link_dir = _link_dir(tmp_path, "mega", link_id)
    assert not (link_dir / ".megatmp.AbCd").exists()
    db_sync.expire_all()
    assert db_sync.get(ExternalLink, link_id).status == "failed"


def test_a_retry_clears_what_a_killed_attempt_left(db_sync, tmp_path, monkeypatch):
    """Leftovers from an attempt whose worker was killed are removed before the
    next fetch starts; an image imported in place by an earlier attempt stays."""
    _, link = _seed(db_sync, host="pixeldrain")
    link_id = link.id
    monkeypatch.setattr(ext, "IMAGES_ROOT", tmp_path)
    monkeypatch.setattr(ext, "_redis", lambda: _FakeRedis())
    link_dir = _link_dir(tmp_path, "pixeldrain", link_id)
    link_dir.mkdir(parents=True)
    leftover = link_dir / "build.zip.part"
    leftover.write_bytes(b"partial")
    imported = link_dir / "kept.jpg"
    imported.write_bytes(_structured_jpeg_bytes())
    db_sync.add(ImageRecord(
        path=str(imported), sha256="ext-kept".ljust(64, "x"),
        size_bytes=imported.stat().st_size, mime="image/jpeg",
        width=256, height=256, origin="downloaded", integrity_status="ok",
    ))
    db_sync.commit()
    seen_at_fetch = []

    def fake_fetch(host, url, dest_dir, **kwargs):
        seen_at_fetch.extend(sorted(p.name for p in dest_dir.iterdir()))
        return FetchResult(error="HTTP 503")

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

    assert seen_at_fetch == ["kept.jpg"]
    assert imported.exists()
    assert not leftover.exists()


def test_a_rejected_download_is_removed(db_sync, tmp_path, monkeypatch):
    """A downloaded image the import filters reject has no record, so it is
    removed rather than left in the link's dir."""
    from backend.app.models import ImportSettings

    _, link = _seed(db_sync, host="dropbox")
    link_id = link.id
    settings = db_sync.execute(
        select(ImportSettings).where(ImportSettings.id == 1)
    ).scalar_one()
    settings.min_width = 512
    db_sync.commit()
    monkeypatch.setattr(ext, "IMAGES_ROOT", tmp_path)
    monkeypatch.setattr(ext, "_redis", lambda: _FakeRedis())
    written = []

    def fake_fetch(host, url, dest_dir, **kwargs):
        dest_dir.mkdir(parents=True, exist_ok=True)
        f = dest_dir / "small.jpg"  # 256px wide, below min_width
        f.write_bytes(_structured_jpeg_bytes())
        written.append(f)
        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["images"] == 0
    assert not written[0].exists()
    db_sync.expire_all()
    assert db_sync.get(ExternalLink, link_id).status == "downloaded"


def test_sweep_recovers_a_link_whose_worker_died(db_sync, monkeypatch):
    """A link 'downloading' past the task's hard limit (or with no claim time)
    counts as a failed attempt and is retried; a fresh claim is left alone."""
    from datetime import UTC, datetime, timedelta

    _, stale = _seed(db_sync, host="mega", status="downloading")
    _, unclaimed = _seed(db_sync, host="gdrive", status="downloading")
    _, fresh = _seed(db_sync, host="dropbox", status="downloading")
    now = datetime.now(UTC)
    stale.started_at = now - timedelta(minutes=ext._STALL_MINUTES + 5)
    fresh.started_at = now - timedelta(minutes=5)
    db_sync.commit()
    stale_id, unclaimed_id, fresh_id = stale.id, unclaimed.id, fresh.id

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

    assert out["recovered"] == 2
    assert set(enqueued) == {stale_id, unclaimed_id}
    db_sync.expire_all()
    for lid in (stale_id, unclaimed_id):
        row = db_sync.get(ExternalLink, lid)
        assert (row.status, row.attempts) == ("failed", 1)
    assert db_sync.get(ExternalLink, fresh_id).status == "downloading"