stack/tests/prisma/test_fetch_deep.py

"""Deeper tests for prisma.fetch — run() orchestrator + attach_pdf."""

from __future__ import annotations

import threading
import time

import httpx

from prisma import fetch as fetch_mod
from prisma.fetch import PendingItem, attach_pdf, fetch_one, pending_queue, run
from zot.db import TYPE_MAP, Db
from zot.schema import create_db


class TestAttachPdf:
    def test_creates_attachment(self, tmp_path):
        path = str(tmp_path / "z.sqlite")
        con = create_db(path)
        con.close()
        with Db(path) as db:
            iid = db.create_item(TYPE_MAP["document"])
            db.commit()
            pdf = tmp_path / "test.pdf"
            pdf.write_bytes(b"%PDF-1.7 test content")
            storage = tmp_path / "storage"
            storage.mkdir()
            att_id = attach_pdf(db, iid, pdf, storage, title="test.pdf")
            assert att_id > 0


class TestPendingQueue:
    def test_empty_queue(self, tmp_path):
        path = str(tmp_path / "z.sqlite")
        con = create_db(path)
        con.close()
        with Db(path) as db:
            result = pending_queue(db, "nonexistent")
            assert result == []


class TestFallbackSerialized:
    def test_fallback_tier_never_concurrent(self, tmp_path, monkeypatch):
        """Parallel fetch_one calls must enter the fallback tier one at a time."""
        state = {"active": 0, "max_active": 0}
        gauge = threading.Lock()

        def fake_fallback(client, doi):
            with gauge:
                state["active"] += 1
                state["max_active"] = max(state["max_active"], state["active"])
            time.sleep(0.05)
            with gauge:
                state["active"] -= 1
            return None  # miss → cascade ends, no download

        # direct tiers all miss so every call reaches the fallback tier
        monkeypatch.setattr(fetch_mod, "fetch_unpaywall", lambda c, d, e: None)
        monkeypatch.setattr(fetch_mod, "fetch_pmc", lambda c, p: None)
        monkeypatch.setattr(fetch_mod, "fetch_fallback", fake_fallback)

        client = httpx.Client(
            transport=httpx.MockTransport(lambda r: httpx.Response(404))
        )
        items = [
            PendingItem(zot_id=i, doi=f"10.1/{i}", pmid="", pmcid="", title=f"t{i}")
            for i in range(6)
        ]
        threads = [
            threading.Thread(
                target=fetch_one,
                args=(it,),
                kwargs={
                    "email": "e@x.com",
                    "scratch": tmp_path,
                    "client": client,
                    "client_proxied": client,
                },
            )
            for it in items
        ]
        for t in threads:
            t.start()
        for t in threads:
            t.join()
        assert state["max_active"] == 1


def _seed_project(path, n):
    """Seed *n* screen:include items in project:test into an existing DB."""
    ids = []
    with Db(path) as db:
        for i in range(n):
            item_id = db.create_item(TYPE_MAP["journalArticle"])
            db.set_fields(item_id, {"title": f"Article {i}", "DOI": f"10.1/{i}"})
            db.tag_item(item_id, "project:test")
            db.tag_item(item_id, "screen:include")
            ids.append(item_id)
        db.commit()
    return ids


class TestRunParallel:
    def test_all_items_attached_and_counted(self, zotero_db, tmp_path, monkeypatch):
        _seed_project(zotero_db, 6)

        def fake_fetch_one(item, *, email, scratch, client, client_proxied=None):
            time.sleep(0.05)
            p = scratch / f"{item.zot_id}.pdf"
            p.write_bytes(b"%PDF-1.4 fake")
            return "unpaywall", p

        monkeypatch.setattr(fetch_mod, "fetch_one", fake_fetch_one)
        with Db(zotero_db) as db:
            stats = fetch_mod.run(
                db,
                project="test",
                storage_dir=tmp_path / "storage",
                scratch_dir=tmp_path / "scratch",
                email="e@x.com",
                fetch_proxy=None,
                workers=4,
            )
            assert stats["unpaywall"] == 6
            assert stats["missed"] == 0
            assert stats["errors"] == 0
            # every item now has a PDF attachment → queue drains
            assert fetch_mod.pending_queue(db, "test") == []

    def test_fetches_overlap(self, zotero_db, tmp_path, monkeypatch):
        _seed_project(zotero_db, 6)
        state = {"active": 0, "max_active": 0}
        gauge = threading.Lock()

        def slow_fetch_one(item, *, email, scratch, client, client_proxied=None):
            with gauge:
                state["active"] += 1
                state["max_active"] = max(state["max_active"], state["active"])
            time.sleep(0.2)
            with gauge:
                state["active"] -= 1
            return None

        monkeypatch.setattr(fetch_mod, "fetch_one", slow_fetch_one)
        with Db(zotero_db) as db:
            stats = fetch_mod.run(
                db,
                project="test",
                storage_dir=tmp_path / "storage",
                scratch_dir=tmp_path / "scratch",
                email="e@x.com",
                fetch_proxy=None,
                workers=6,
            )
        assert stats["missed"] == 6
        assert state["max_active"] >= 2

    def test_worker_exception_counts_as_error(self, zotero_db, tmp_path, monkeypatch):
        _seed_project(zotero_db, 3)

        def boom(item, *, email, scratch, client, client_proxied=None):
            raise RuntimeError("kaput")

        monkeypatch.setattr(fetch_mod, "fetch_one", boom)
        with Db(zotero_db) as db:
            stats = fetch_mod.run(
                db,
                project="test",
                storage_dir=tmp_path / "storage",
                scratch_dir=tmp_path / "scratch",
                email="e@x.com",
                fetch_proxy=None,
                workers=2,
            )
        assert stats["errors"] == 3


class TestRun:
    def test_empty_queue(self, tmp_path):
        path = str(tmp_path / "z.sqlite")
        con = create_db(path)
        con.close()
        with Db(path) as db:
            stats = run(
                db,
                project="test",
                storage_dir=tmp_path / "storage",
                scratch_dir=tmp_path / "scratch",
                email="test@test.com",
                fetch_proxy=None,
            )
            assert stats["missed"] == 0
            assert stats["unpaywall"] == 0