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 _other_jpeg_bytes():
    # A second structured image, unlike _structured_jpeg_bytes, so the two
    # don't near-duplicate each other at import.
    im = Image.new("RGB", (256, 256))
    px = im.load()
    for y in range(256):
        for x in range(256):
            px[x, y] = ((x * x) % 256, (y * 3) % 256, (x ^ y) % 256)
    buf = io.BytesIO()
    im.save(buf, "JPEG")
    return buf.getvalue()


def test_a_failed_fetch_keeps_complete_files(db_sync, tmp_path, monkeypatch):
    """A failed fetch removes staging files only. A complete file may be the
    only copy left once the host has dropped the upload."""
    _, 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 / "images").mkdir(parents=True, exist_ok=True)
        (dest_dir / "images" / "done.jpg").write_bytes(_structured_jpeg_bytes())
        (dest_dir / "images" / ".megatmp.AbCd").write_bytes(b"half")
        return FetchResult(error="API call 'f' failed: ENOENT")

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

    link_dir = _link_dir(tmp_path, "mega", link_id)
    assert (link_dir / "images" / "done.jpg").exists()
    assert not (link_dir / "images" / ".megatmp.AbCd").exists()


def test_a_successful_fetch_imports_what_an_earlier_attempt_finished(
    db_sync, tmp_path, monkeypatch,
):
    """Complete files an earlier attempt wrote, in subfolders too, are imported
    with the new fetch's files instead of being deleted or left behind."""
    _, link = _seed(db_sync, host="mega")
    link_id = link.id
    monkeypatch.setattr(ext, "IMAGES_ROOT", tmp_path)
    monkeypatch.setattr(ext, "_redis", lambda: _FakeRedis())
    link_dir = _link_dir(tmp_path, "mega", link_id)
    (link_dir / "pack").mkdir(parents=True)
    earlier = link_dir / "pack" / "earlier.jpg"
    earlier.write_bytes(_structured_jpeg_bytes())

    def fake_fetch(host, url, dest_dir, **kwargs):
        f = dest_dir / "new.jpg"
        f.write_bytes(_other_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"] == 2
    db_sync.expire_all()
    paths = set(db_sync.execute(select(ImageRecord.path)).scalars())
    assert {str(earlier), str(link_dir / "new.jpg")} <= paths


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"


def test_recovery_gives_a_dead_link_one_more_attempt_at_growing_intervals(db_sync, monkeypatch):
    """A dead link is retried once after 2**recoveries days, not given a fresh
    three-attempt budget, so a link that always fails isn't downloaded daily."""
    from datetime import UTC, datetime, timedelta

    now = datetime.now(UTC)
    _, due = _seed(db_sync, host="mega", status="dead")
    _, waiting = _seed(db_sync, host="gdrive", status="dead")
    _, spent = _seed(db_sync, host="dropbox", status="dead")
    due.attempts = waiting.attempts = spent.attempts = ext.DEAD_LETTER_THRESHOLD
    due.recoveries, due.completed_at = 2, now - timedelta(days=4, hours=1)
    waiting.recoveries, waiting.completed_at = 2, now - timedelta(days=3)
    spent.recoveries, spent.completed_at = ext._MAX_RECOVERIES, now - timedelta(days=90)
    db_sync.commit()
    due_id, waiting_id, spent_id = due.id, waiting.id, spent.id

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

    assert out["revived"] == 1
    assert enqueued == [due_id]
    db_sync.expire_all()
    row = db_sync.get(ExternalLink, due_id)
    assert (row.status, row.attempts, row.recoveries) == (
        "failed", ext.DEAD_LETTER_THRESHOLD - 1, 3,
    )
    assert db_sync.get(ExternalLink, waiting_id).status == "dead"
    assert db_sync.get(ExternalLink, spent_id).status == "dead"


def test_a_recovered_link_that_fails_again_goes_straight_back_to_dead(
    db_sync, tmp_path, monkeypatch,
):
    _, link = _seed(db_sync, host="mega", status="dead")
    link.attempts = ext.DEAD_LETTER_THRESHOLD
    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="too big"))
    monkeypatch.setattr(ext.fetch_external_link, "delay", ext.fetch_external_link)

    ext.recover_external_links()

    db_sync.expire_all()
    row = db_sync.get(ExternalLink, link_id)
    assert (row.status, row.recoveries) == ("dead", 1)


def test_prune_removes_only_links_recovery_has_given_up_on(db_sync):
    from datetime import UTC, datetime, timedelta

    old = datetime.now(UTC) - timedelta(days=ext._RETENTION_DAYS + 1)
    _, spent = _seed(db_sync, host="mega", status="dead")
    _, recovering = _seed(db_sync, host="gdrive", status="dead")
    _, recent = _seed(db_sync, host="dropbox", status="dead")
    spent.recoveries, spent.completed_at = ext._MAX_RECOVERIES, old
    recovering.recoveries, recovering.completed_at = 1, old
    recent.recoveries = ext._MAX_RECOVERIES
    recent.completed_at = datetime.now(UTC) - timedelta(days=1)
    db_sync.commit()
    spent_id, recovering_id, recent_id = spent.id, recovering.id, recent.id

    assert ext.prune_external_links() == {"pruned": 1}
    db_sync.expire_all()
    assert db_sync.get(ExternalLink, spent_id) is None
    assert db_sync.get(ExternalLink, recovering_id) is not None
    assert db_sync.get(ExternalLink, recent_id) is not None