stack/tests/nlm/test_vsac.py

"""nlm.vsac — paged definitions, paged expansions, resumable hoard."""

from __future__ import annotations

import gzip
import json

import httpx
import pytest

from nlm import vsac

KEY = "0123-secret-key"


def _client(handler) -> httpx.Client:
    return httpx.Client(base_url=vsac.FHIR_URL, transport=httpx.MockTransport(handler))


def _vs(i: int) -> dict:
    return {"resourceType": "ValueSet", "id": f"1.2.{i}-20260101", "title": f"VS {i}"}


def _read_gz(path):
    with gzip.open(path, "rt", encoding="utf-8") as f:
        return json.load(f)


def test_make_client_uses_basic_auth_and_fhir_json():
    c = vsac.make_client(KEY)
    try:
        assert str(c.base_url).rstrip("/") == vsac.FHIR_URL
        assert c.headers["Accept"] == "application/fhir+json"
        assert isinstance(c.auth, httpx.BasicAuth)
    finally:
        c.close()


def test_get_backs_off_on_throttling_then_succeeds():
    calls, slept = [], []

    def handler(request):
        calls.append(request)
        if len(calls) == 1:
            return httpx.Response(429, headers={"Retry-After": "3"})
        if len(calls) == 2:
            return httpx.Response(503)
        return httpx.Response(200, json={"ok": True})

    with _client(handler) as c:
        assert vsac._get(c, "/x", sleep=slept.append) == {"ok": True}
    assert slept == [3.0, 2.0]


def test_get_gives_up_after_the_last_try():
    with _client(lambda r: httpx.Response(503)) as c:
        with pytest.raises(httpx.HTTPStatusError):
            vsac._get(c, "/x", tries=2, sleep=lambda s: None)


def test_iter_definitions_follows_next_links_over_https():
    pages = {
        0: {
            "entry": [{"resource": _vs(1)}, {"resource": _vs(2)}],
            "link": [
                {
                    "relation": "next",
                    "url": "http://cts.nlm.nih.gov/fhir/res/ValueSet?_offset=2",
                }
            ],
        },
        2: {
            "entry": [{"resource": _vs(3)}],
            "link": [{"relation": "self", "url": "x"}],
        },
    }

    def handler(request):
        assert request.url.scheme == "https"
        return httpx.Response(
            200, json=pages[int(request.url.params.get("_offset", 0))]
        )

    with _client(handler) as c:
        ids = [vs["id"] for vs in vsac.iter_definitions(c, page_size=2)]
    assert ids == ["1.2.1-20260101", "1.2.2-20260101", "1.2.3-20260101"]


def test_save_definitions_writes_files_and_index_and_keeps_existing(tmp_path):
    existing = tmp_path / "definitions" / "1.2.1-20260101.json.gz"
    existing.parent.mkdir(parents=True)
    with gzip.open(existing, "wt") as f:
        json.dump({"id": "kept"}, f)

    with _client(
        lambda r: httpx.Response(
            200, json={"entry": [{"resource": _vs(1)}, {"resource": _vs(2)}]}
        )
    ) as c:
        n = vsac.save_definitions(c, tmp_path)

    assert n == 2
    assert _read_gz(existing) == {"id": "kept"}
    assert (
        _read_gz(tmp_path / "definitions" / "1.2.2-20260101.json.gz")["title"] == "VS 2"
    )
    index = [
        json.loads(line) for line in (tmp_path / "index.jsonl").read_text().splitlines()
    ]
    assert [row["id"] for row in index] == ["1.2.1-20260101", "1.2.2-20260101"]
    assert index[0]["status"] == ""


def _expanding(codes: list[str], *, total: bool = True, stall_at: int | None = None):
    def handler(request):
        offset = int(request.url.params["offset"])
        count = int(request.url.params["count"])
        page = (
            []
            if stall_at is not None and offset >= stall_at
            else codes[offset : offset + count]
        )
        exp = {"contains": [{"system": "s", "code": c} for c in page], "offset": offset}
        if total:
            exp["total"] = len(codes)
        return httpx.Response(200, json={"resourceType": "ValueSet", "expansion": exp})

    return handler


def test_expand_joins_pages():
    codes = [str(i) for i in range(7)]
    with _client(_expanding(codes)) as c:
        out = vsac.expand(c, "1.2.3", page=3)
    assert [x["code"] for x in out["expansion"]["contains"]] == codes
    assert "offset" not in out["expansion"]


def test_expand_stops_when_a_page_comes_back_empty():
    with _client(_expanding([str(i) for i in range(7)], stall_at=3)) as c:
        out = vsac.expand(c, "1.2.3", page=3)
    assert len(out["expansion"]["contains"]) == 3


def test_expand_without_total_takes_one_page():
    with _client(_expanding(["a", "b"], total=False)) as c:
        out = vsac.expand(c, "1.2.3")
    assert [x["code"] for x in out["expansion"]["contains"]] == ["a", "b"]


def test_hoard_skips_present_expands_new_and_logs_failures(tmp_path):
    present = tmp_path / "expansions" / "done.json.gz"
    present.parent.mkdir(parents=True)
    present.write_bytes(b"")

    def handler(request):
        if "/ValueSet/bad/" in request.url.path:
            return httpx.Response(404, text=f"no such value set apiKey={KEY}")
        return httpx.Response(
            200, json={"expansion": {"total": 1, "contains": [{"code": "x"}]}}
        )

    made, ticks = [], []

    def factory():
        made.append(1)
        return _client(handler)

    stats = vsac.hoard(
        factory,
        tmp_path,
        ["done", "good", "bad"],
        api_key=KEY,
        workers=2,
        progress=lambda done, total: ticks.append((done, total)),
    )
    assert stats == {"expanded": 1, "present": 1, "failed": 1}
    assert _read_gz(tmp_path / "expansions" / "good.json.gz")["expansion"][
        "contains"
    ] == [{"code": "x"}]
    err = [
        json.loads(line)
        for line in (tmp_path / "errors.jsonl").read_text().splitlines()
    ]
    assert err[0]["id"] == "bad" and KEY not in err[0]["error"]
    assert sorted(ticks) == [(1, 2), (2, 2)]
    assert 1 <= len(made) <= 2


def test_hoard_without_progress_callback(tmp_path):
    def handler(request):
        return httpx.Response(200, json={"expansion": {"contains": []}})

    stats = vsac.hoard(lambda: _client(handler), tmp_path, ["a"], workers=1)
    assert stats["expanded"] == 1