stack/tests/zot/fetch/test_run.py

from __future__ import annotations

from pathlib import Path

import httpx
import pytest

from tests.zot.fetch.conftest import mk_item
from zot.db import Db
from zot.fetch.attach import has_full_text
from zot.fetch.deliver import MIN_BYTES, AnnasAuthError, Budget
from zot.fetch.ledger import Ledger
from zot.fetch.resolve import BROWSER_UA
from zot.fetch.run import STATS, Clients, make_clients, run, safe_name
from zot.fetch.select import select_items

FX = Path(__file__).parent / "fixtures"
PDF = b"%PDF-1.4\n" + b"0" * MIN_BYTES
PDF_HDR = {"content-type": "application/pdf"}
ADS = (FX / "libgen_ads.html").read_text()
DOI = "10.1056/NEJMoa1000678"


def _py_copier(src: Path, dest_dir: Path) -> None:
    dest_dir.mkdir(parents=True, exist_ok=True)
    (dest_dir / src.name).write_bytes(src.read_bytes())


def _clients(direct_handler, proxied_handler) -> Clients:
    return Clients(
        direct=httpx.Client(transport=httpx.MockTransport(direct_handler)),
        proxied=httpx.Client(
            transport=httpx.MockTransport(proxied_handler),
            headers={"User-Agent": BROWSER_UA},
        ),
    )


def _unpaywall_ok(r: httpx.Request) -> httpx.Response:
    if r.url.host == "api.unpaywall.org":
        return httpx.Response(200, text=(FX / "unpaywall_doi.json").read_text())
    return httpx.Response(200, content=PDF, headers=PDF_HDR)


def _unpaywall_miss(r: httpx.Request) -> httpx.Response:
    return httpx.Response(404)


def _libgen_ok(r: httpx.Request) -> httpx.Response:
    q = r.url.params
    if r.url.path == "/json.php" and q.get("object") == "e":
        return httpx.Response(200, text=(FX / "libgen_json_doi.json").read_text())
    if r.url.path == "/json.php":
        return httpx.Response(200, text=(FX / "libgen_json_file.json").read_text())
    if r.url.path == "/ads.php":
        return httpx.Response(200, text=ADS)
    if r.url.path == "/get.php":
        return httpx.Response(200, content=PDF, headers=PDF_HDR)
    return httpx.Response(404)


def _libgen_none(r: httpx.Request) -> httpx.Response:
    return httpx.Response(200, text="[]")


def _libgen_dead(r: httpx.Request) -> httpx.Response:
    """Every route answers definitively (a miss), nothing errors."""
    if r.url.path == "/json.php" and r.url.params.get("object") == "e":
        return httpx.Response(200, text=(FX / "libgen_json_doi.json").read_text())
    if r.url.path == "/json.php":
        return httpx.Response(200, text=(FX / "libgen_json_file.json").read_text())
    if r.url.path == "/ads.php":
        return httpx.Response(200, text=ADS)
    if r.url.host == "annas-archive.gl":
        return httpx.Response(200, json={"download_url": None, "error": "nope"})
    return httpx.Response(404)


def _setup(tmp_path: Path, **kw):
    return dict(
        ledger=Ledger(tmp_path / "ledger.jsonl"),
        budget=kw.pop("budget", Budget()),
        email="dev@fhirworx.io",
        storage_dir=tmp_path / "storage",
        scratch_dir=tmp_path / "scratch",
        copier=kw.pop("copier", _py_copier),
        sleep=lambda s: None,
        out=kw.pop("out", lambda *a: None),
        **kw,
    )


def _zero(**over) -> dict[str, int]:
    d = dict.fromkeys(STATS, 0)
    d.update(over)
    return d


def test_safe_name():
    assert safe_name("Early Palliative Care: NSCLC / 2010", "K") == (
        "Early Palliative Care NSCLC 2010"
    )
    assert safe_name("", "ABCD2345") == "ABCD2345"
    assert len(safe_name("x" * 200, "K")) == 80


def test_make_clients_sets_proxy_and_agents():
    c = make_clients("http://127.0.0.1:19035", email="dev@fhirworx.io")
    try:
        assert c.proxied.headers["user-agent"] == BROWSER_UA
        assert "dev@fhirworx.io" in c.direct.headers["user-agent"]
        assert c.proxied._transport is not c.direct._transport
    finally:
        c.direct.close()
        c.proxied.close()


def test_run_attaches_from_unpaywall_and_ledgers(zdb: Db, tmp_path: Path):
    item = mk_item(zdb, title="Early Palliative Care", doi=DOI)
    refs = select_items(zdb, all_=True)
    lines: list[str] = []
    kw = _setup(tmp_path, out=lambda *a: lines.append(" ".join(map(str, a))))
    stats = run(zdb, refs, clients=_clients(_unpaywall_ok, _libgen_none), **kw)
    assert stats == _zero(seen=1, attempted=1, attached=1)
    assert has_full_text(zdb, item)
    assert "source:unpaywall" in zdb.get_tags(item)
    att = zdb.con.execute(
        "SELECT path FROM itemAttachments WHERE parentItemID = ?", (item,)
    ).fetchone()[0]
    assert att == "storage:Early Palliative Care.pdf"
    assert kw["ledger"].summary() == {"attached": 1}
    assert any("attached" in ln for ln in lines)
    assert not list((tmp_path / "scratch").iterdir())  # scratch file cleaned up


def test_run_falls_through_to_libgen_with_md5_provenance(zdb: Db, tmp_path: Path):
    item = mk_item(zdb, title="T", doi=DOI)
    kw = _setup(tmp_path)
    refs = select_items(zdb, all_=True)
    stats = run(zdb, refs, clients=_clients(_unpaywall_miss, _libgen_ok), **kw)
    assert stats["attached"] == 1
    assert "source:libgen" in zdb.get_tags(item)
    assert zdb.get_field(item, "extra") == "md5: a78add1bd62047ead5945134488f470f"


def test_run_no_candidate_and_all_failed_outcomes(zdb: Db, tmp_path: Path):
    mk_item(zdb, title="no ids")
    mk_item(zdb, title="dead", doi=DOI)
    kw = _setup(tmp_path)
    refs = select_items(zdb, all_=True)
    stats = run(zdb, refs, clients=_clients(_unpaywall_miss, _libgen_dead), **kw)
    assert stats == _zero(seen=2, attempted=2, no_candidate=1, all_failed=1)
    assert kw["ledger"].summary() == {"no_candidate": 1, "all_failed": 1}


def test_run_ledgers_quota_when_only_annas_was_left(zdb: Db, tmp_path: Path):
    mk_item(zdb, title="dead", doi=DOI)
    kw = _setup(tmp_path, budget=Budget(annas_key="k", annas_left=1, reserve=5))
    refs = select_items(zdb, all_=True)
    stats = run(zdb, refs, clients=_clients(_unpaywall_miss, _libgen_dead), **kw)
    assert stats["quota"] == 1 and stats["all_failed"] == 0
    assert kw["ledger"].summary() == {"quota": 1}


def test_run_skips_ledgered_items_and_honours_limit(zdb: Db, tmp_path: Path):
    a = mk_item(zdb, title="A", doi=DOI)
    mk_item(zdb, title="B", doi=DOI)
    mk_item(zdb, title="C", doi=DOI)
    refs = select_items(zdb, all_=True)
    kw = _setup(tmp_path)
    kw["ledger"].record(refs[0].key, outcome="no_candidate")
    clients = _clients(_unpaywall_ok, _libgen_none)
    stats = run(zdb, refs, clients=clients, limit=1, **kw)
    assert stats == _zero(seen=2, skipped=1, attempted=1, attached=1)
    assert not has_full_text(zdb, a)


def test_run_retry_flag_reprocesses_ledgered_items(zdb: Db, tmp_path: Path):
    mk_item(zdb, title="A", doi=DOI)
    refs = select_items(zdb, all_=True)
    kw = _setup(tmp_path)
    kw["ledger"].record(refs[0].key, outcome="all_failed")
    clients = _clients(_unpaywall_ok, _libgen_none)
    stats = run(zdb, refs, clients=clients, retry=True, **kw)
    assert stats["skipped"] == 0 and stats["attached"] == 1


def test_run_dry_run_prints_candidates_and_writes_nothing(zdb: Db, tmp_path: Path):
    item = mk_item(zdb, title="T", doi=DOI)
    mk_item(zdb, title="no ids")
    lines: list[str] = []
    kw = _setup(tmp_path, out=lambda *a: lines.append(" ".join(map(str, a))))
    refs = select_items(zdb, all_=True)
    clients = _clients(_unpaywall_ok, _libgen_ok)
    stats = run(zdb, refs, clients=clients, dry_run=True, **kw)
    assert stats == _zero(seen=2, attempted=2)
    assert not has_full_text(zdb, item)
    assert not (tmp_path / "ledger.jsonl").exists()
    assert any("unpaywall" in ln for ln in lines)
    assert any("libgen" in ln for ln in lines)
    assert any("no candidates" in ln for ln in lines)


def test_run_annas_auth_error_aborts(zdb: Db, tmp_path: Path):
    mk_item(zdb, title="T", doi=DOI)

    def proxied(r: httpx.Request) -> httpx.Response:
        if r.url.host == "annas-archive.gl":
            return httpx.Response(401)
        return _libgen_dead(r)

    kw = _setup(tmp_path, budget=Budget(annas_key="bad"))
    refs = select_items(zdb, all_=True)
    with pytest.raises(AnnasAuthError):
        run(zdb, refs, clients=_clients(_unpaywall_miss, proxied), **kw)


# ── fix pass (review) ─────────────────────────────────────────────


def _proxy_down(r: httpx.Request) -> httpx.Response:
    raise httpx.ConnectError("proxy down")


def test_run_transient_failure_is_error_not_terminal(zdb: Db, tmp_path: Path):
    mk_item(zdb, title="T", doi=DOI)
    kw = _setup(tmp_path)
    refs = select_items(zdb, all_=True)
    stats = run(zdb, refs, clients=_clients(_unpaywall_miss, _proxy_down), **kw)
    assert stats == _zero(seen=1, attempted=1, error=1)
    led = kw["ledger"]
    assert led.summary() == {"error": 1}
    assert led.should_try(refs[0].key) is True  # retried next run, not 30 days out
    line = (tmp_path / "ledger.jsonl").read_text()
    assert "ConnectError" in line


def test_run_all_failed_and_no_candidate_carry_notes(zdb: Db, tmp_path: Path):
    mk_item(zdb, title="no ids")
    mk_item(zdb, title="dead", doi=DOI)
    kw = _setup(tmp_path)
    refs = select_items(zdb, all_=True)
    run(zdb, refs, clients=_clients(_unpaywall_miss, _libgen_dead), **kw)
    lines = (tmp_path / "ledger.jsonl").read_text().splitlines()
    assert "no identifiers" in lines[0]
    assert "unpaywall" in lines[1] and "libgen get: http 404" in lines[1]


def test_run_isolates_a_failing_item_and_rolls_back(zdb: Db, tmp_path: Path):
    bad = mk_item(zdb, title="bad", doi=DOI)
    good = mk_item(zdb, title="good", doi=DOI)
    calls = {"n": 0}

    def flaky_copier(src: Path, dest_dir: Path) -> None:
        calls["n"] += 1
        if calls["n"] == 1:
            raise RuntimeError("sudo: a password is required")
        _py_copier(src, dest_dir)

    kw = _setup(tmp_path, copier=flaky_copier)
    refs = select_items(zdb, all_=True)
    stats = run(zdb, refs, clients=_clients(_unpaywall_ok, _libgen_none), **kw)
    assert stats == _zero(seen=2, attempted=2, attached=1, error=1)
    # the failed item left no rows behind (rolled back) and no scratch file
    rows = zdb.con.execute(
        "SELECT COUNT(*) FROM itemAttachments WHERE parentItemID = ?", (bad,)
    ).fetchone()[0]
    assert rows == 0 and "source:unpaywall" not in zdb.get_tags(bad)
    assert has_full_text(zdb, good)
    assert not list((tmp_path / "scratch").iterdir())
    assert "RuntimeError" in (tmp_path / "ledger.jsonl").read_text()
    assert kw["ledger"].should_try(refs[0].key) is True


def test_run_annas_auth_error_still_aborts_after_fix(zdb: Db, tmp_path: Path):
    mk_item(zdb, title="T", doi=DOI)

    def proxied(r: httpx.Request) -> httpx.Response:
        if r.url.host == "annas-archive.gl":
            return httpx.Response(403, json={"error": "Invalid secret key"})
        return _libgen_dead(r)

    kw = _setup(tmp_path, budget=Budget(annas_key="bad"))
    refs = select_items(zdb, all_=True)
    with pytest.raises(AnnasAuthError):
        run(zdb, refs, clients=_clients(_unpaywall_miss, proxied), **kw)


def test_run_delivery_error_is_error_outcome(zdb: Db, tmp_path: Path):
    """Candidates exist but every delivery route errored (5xx) → retry next run."""

    def flaky(r: httpx.Request) -> httpx.Response:
        if r.url.path == "/ads.php":
            return httpx.Response(503)
        return _libgen_ok(r)

    mk_item(zdb, title="T", doi=DOI)
    kw = _setup(tmp_path)
    refs = select_items(zdb, all_=True)
    stats = run(zdb, refs, clients=_clients(_unpaywall_miss, flaky), **kw)
    assert stats == _zero(seen=1, attempted=1, error=1)
    assert "503" in (tmp_path / "ledger.jsonl").read_text()