stack/tests/zot/fetch/test_deliver.py

from __future__ import annotations

from pathlib import Path

import httpx
import pytest

from zot.fetch.deliver import (
    MIN_BYTES,
    AnnasAuthError,
    Budget,
    download,
    fetch,
    fetch_annas,
    fetch_libgen,
)
from zot.fetch.resolve import LIBGEN, Candidate, Pacer

FX = Path(__file__).parent / "fixtures"
PDF = b"%PDF-1.4\n" + b"0" * MIN_BYTES
EPUB = b"PK\x03\x04" + b"0" * MIN_BYTES
MD5 = "a78add1bd62047ead5945134488f470f"
ADS = (FX / "libgen_ads.html").read_text()
ANNAS_OK = (FX / "annas_fast_download.json").read_text()
PDF_HDR = {"content-type": "application/pdf"}


def _pacer() -> Pacer:
    return Pacer(3.0, clock=lambda: 0.0, sleep=lambda s: None)


def _client(handler) -> httpx.Client:
    return httpx.Client(transport=httpx.MockTransport(handler))


def _lg(md5=MD5) -> Candidate:
    return Candidate("libgen", md5, None, "pdf", None, "T")


def _up() -> Candidate:
    return Candidate("unpaywall", None, "https://example.org/x.pdf", "pdf", None, "T")


def _nosleep(s: float) -> None:
    return None


# ── download ──────────────────────────────────────────────────────


@pytest.mark.parametrize(
    "status, ctype, body, ext, ok",
    [
        (200, "application/pdf", PDF, "pdf", True),
        (200, "application/octet-stream", PDF, "pdf", True),
        (200, "application/epub+zip", EPUB, "epub", True),
        (200, "text/html; charset=utf-8", PDF, "pdf", False),
        (200, "application/pdf", b"%PDF tiny", "pdf", False),
        (200, "application/pdf", b"<html>" + b"0" * MIN_BYTES, "pdf", False),
        (503, "application/pdf", PDF, "pdf", False),
    ],
)
def test_download_validates_status_type_size_and_magic(
    tmp_path, status, ctype, body, ext, ok
):
    c = _client(
        lambda r: httpx.Response(status, content=body, headers={"content-type": ctype})
    )
    dest = tmp_path / f"f.{ext}"
    got = download(c, "https://x/y", dest, ext=ext)
    assert (got == dest) is ok
    assert dest.exists() is ok


def test_download_transport_error_is_none(tmp_path):
    def boom(r):
        raise httpx.ReadTimeout("t")

    assert download(_client(boom), "https://x", tmp_path / "f.pdf", ext="pdf") is None


def test_download_passes_headers(tmp_path):
    seen = {}

    def h(r):
        seen.update(r.headers)
        return httpx.Response(200, content=PDF, headers=PDF_HDR)

    download(
        _client(h),
        "https://x",
        tmp_path / "f.pdf",
        ext="pdf",
        headers={"Referer": "https://r/"},
    )
    assert seen["referer"] == "https://r/"


# ── libgen ────────────────────────────────────────────────────────


def _libgen_handler(ads_pages: list[httpx.Response], get_status=200):
    calls: list[str] = []

    def h(r: httpx.Request) -> httpx.Response:
        calls.append(r.url.path)
        if r.url.path == "/ads.php":
            assert r.url.params["md5"] == MD5
            assert r.headers["referer"] == f"{LIBGEN}/"
            return ads_pages.pop(0)
        if r.url.path == "/get.php":
            assert r.url.params["key"] == "AJPW7IO3P1LGRGWA"
            return httpx.Response(get_status, content=PDF, headers=PDF_HDR)
        raise AssertionError(r.url)

    return h, calls


def test_fetch_libgen_ads_then_get(tmp_path):
    h, calls = _libgen_handler([httpx.Response(200, text=ADS)])
    got = fetch_libgen(
        _client(h), _lg(), tmp_path / "f.pdf", pacer=_pacer(), sleep=_nosleep
    )
    assert got == tmp_path / "f.pdf" and calls == ["/ads.php", "/get.php"]


def test_libgen_ads_error_div_retries_once_then_gives_up(tmp_path):
    err = httpx.Response(
        200,
        text='<div class="alert alert-danger">Could not connect to the database 3306</div>',
    )
    h, calls = _libgen_handler([err, err])
    slept: list[float] = []
    got = fetch_libgen(
        _client(h), _lg(), tmp_path / "f.pdf", pacer=_pacer(), sleep=slept.append
    )
    assert got is None and calls == ["/ads.php", "/ads.php"] and slept == [3.0]


def test_libgen_ads_error_then_success_on_retry(tmp_path):
    h, calls = _libgen_handler([httpx.Response(500), httpx.Response(200, text=ADS)])
    got = fetch_libgen(
        _client(h), _lg(), tmp_path / "f.pdf", pacer=_pacer(), sleep=_nosleep
    )
    assert got is not None and calls == ["/ads.php", "/ads.php", "/get.php"]


def test_libgen_ads_link_for_other_md5_is_ignored(tmp_path):
    other = httpx.Response(200, text=ADS.replace(MD5, "f" * 32))
    h, calls = _libgen_handler([other, other])
    got = fetch_libgen(
        _client(h), _lg(), tmp_path / "f.pdf", pacer=_pacer(), sleep=_nosleep
    )
    assert got is None and calls == ["/ads.php", "/ads.php"]


def test_libgen_get_503_is_none(tmp_path):
    h, _ = _libgen_handler([httpx.Response(200, text=ADS)], get_status=503)
    got = fetch_libgen(
        _client(h), _lg(), tmp_path / "f.pdf", pacer=_pacer(), sleep=_nosleep
    )
    assert got is None


def test_libgen_transport_error_is_none(tmp_path):
    def boom(r):
        raise httpx.ConnectError("x")

    got = fetch_libgen(
        _client(boom), _lg(), tmp_path / "f.pdf", pacer=_pacer(), sleep=_nosleep
    )
    assert got is None


# ── Anna's ────────────────────────────────────────────────────────


def _annas_handler(api: httpx.Response, file_status=200):
    def h(r: httpx.Request) -> httpx.Response:
        if r.url.host == "annas-archive.gl":
            assert r.url.path == "/dyn/api/fast_download.json"
            assert r.url.params["md5"] == MD5 and r.url.params["key"] == "SECRET"
            return api
        return httpx.Response(file_status, content=PDF, headers=PDF_HDR)

    return h


def test_fetch_annas_downloads_and_updates_budget(tmp_path):
    b = Budget(annas_key="SECRET")
    c = _client(_annas_handler(httpx.Response(200, text=ANNAS_OK)))
    got = fetch_annas(c, _lg(), tmp_path / "f.pdf", budget=b)
    assert got == tmp_path / "f.pdf" and b.annas_left == 42


@pytest.mark.parametrize("status", [401, 403])
def test_fetch_annas_auth_errors_stop_the_run(tmp_path, status):
    b = Budget(annas_key="SECRET")
    c = _client(_annas_handler(httpx.Response(status, json={"error": "Not a member"})))
    with pytest.raises(AnnasAuthError):
        fetch_annas(c, _lg(), tmp_path / "f.pdf", budget=b)


def test_fetch_annas_other_error_disables_annas(tmp_path):
    b = Budget(annas_key="SECRET")
    c = _client(_annas_handler(httpx.Response(500)))
    assert fetch_annas(c, _lg(), tmp_path / "f.pdf", budget=b) is None
    assert b.annas_disabled is True and b.can_use_annas() is False


def test_fetch_annas_null_download_url_is_none(tmp_path):
    b = Budget(annas_key="SECRET")
    api = httpx.Response(200, json={"download_url": None, "error": "Invalid md5"})
    assert (
        fetch_annas(_client(_annas_handler(api)), _lg(), tmp_path / "f.pdf", budget=b)
        is None
    )


def test_fetch_annas_bad_json_or_transport_is_none(tmp_path):
    b = Budget(annas_key="SECRET")
    c = _client(_annas_handler(httpx.Response(200, text="nope")))
    assert fetch_annas(c, _lg(), tmp_path / "f.pdf", budget=b) is None

    def boom(r):
        raise httpx.ConnectError("x")

    assert fetch_annas(_client(boom), _lg(), tmp_path / "f.pdf", budget=b) is None


def test_fetch_annas_file_failure_still_reports_budget(tmp_path):
    b = Budget(annas_key="SECRET")
    c = _client(_annas_handler(httpx.Response(200, text=ANNAS_OK), file_status=503))
    got = fetch_annas(c, _lg(), tmp_path / "f.pdf", budget=b)
    assert got is None and b.annas_left == 42  # the API call itself consumed one


def test_budget_reserve_blocks_annas():
    assert Budget(annas_key="").can_use_annas() is False
    assert Budget(annas_key="k").can_use_annas() is True  # unknown budget → allowed
    assert Budget(annas_key="k", annas_left=6, reserve=5).can_use_annas() is True
    assert Budget(annas_key="k", annas_left=5, reserve=5).can_use_annas() is False


# ── fetch (the cascade) ───────────────────────────────────────────


def test_fetch_unpaywall_uses_direct_client(tmp_path):
    direct = _client(lambda r: httpx.Response(200, content=PDF, headers=PDF_HDR))

    def never(r):
        raise AssertionError("proxied client used for Unpaywall")

    src = fetch(
        _up(),
        tmp_path / "f.pdf",
        direct=direct,
        proxied=_client(never),
        budget=Budget(),
        pacer=_pacer(),
    )
    assert src == "unpaywall"


def test_fetch_unpaywall_failure_is_none(tmp_path):
    direct = _client(lambda r: httpx.Response(404))
    src = fetch(
        _up(),
        tmp_path / "f.pdf",
        direct=direct,
        proxied=direct,
        budget=Budget(),
        pacer=_pacer(),
    )
    assert src is None


def test_fetch_libgen_then_annas_on_failure(tmp_path):
    def h(r: httpx.Request) -> httpx.Response:
        if r.url.path == "/ads.php":
            return httpx.Response(500)
        if r.url.host == "annas-archive.gl":
            return httpx.Response(200, text=ANNAS_OK)
        return httpx.Response(200, content=PDF, headers=PDF_HDR)

    b = Budget(annas_key="SECRET")
    src = fetch(
        _lg(),
        tmp_path / "f.pdf",
        direct=_client(h),
        proxied=_client(h),
        budget=b,
        pacer=_pacer(),
        sleep=_nosleep,
    )
    assert src == "annas" and b.annas_left == 42


def test_fetch_libgen_success_skips_annas(tmp_path):
    h, calls = _libgen_handler([httpx.Response(200, text=ADS)])
    b = Budget(annas_key="SECRET")
    src = fetch(
        _lg(),
        tmp_path / "f.pdf",
        direct=_client(h),
        proxied=_client(h),
        budget=b,
        pacer=_pacer(),
        sleep=_nosleep,
    )
    assert src == "libgen" and calls == ["/ads.php", "/get.php"]


def test_fetch_marks_quota_when_annas_blocked_by_reserve(tmp_path):
    proxied = _client(lambda r: httpx.Response(500))
    b = Budget(annas_key="SECRET", annas_left=2, reserve=5)
    src = fetch(
        _lg(),
        tmp_path / "f.pdf",
        direct=proxied,
        proxied=proxied,
        budget=b,
        pacer=_pacer(),
        sleep=_nosleep,
    )
    assert src is None and b.quota_hit is True


def test_fetch_all_fail_without_annas_is_not_quota(tmp_path):
    proxied = _client(lambda r: httpx.Response(500))
    b = Budget()
    src = fetch(
        _lg(),
        tmp_path / "f.pdf",
        direct=proxied,
        proxied=proxied,
        budget=b,
        pacer=_pacer(),
        sleep=_nosleep,
    )
    assert src is None and b.quota_hit is False