stack/tests/nlm/test_releases.py

"""nlm.releases — UTS release listing and resumable downloads."""

from __future__ import annotations

import io
import json
import zipfile

import httpx
import pytest

from nlm import scrub
from nlm.releases import Release, current, fetch, release_types, target

KEY = "0123-secret-key"


def _zip_bytes() -> bytes:
    buf = io.BytesIO()
    with zipfile.ZipFile(buf, "w") as z:
        z.writestr("readme.txt", "hello")
    return buf.getvalue()


def _rel(**kw) -> Release:
    base = dict(
        product="SNOMED CT",
        stream="snomed-ct-us-edition",
        name="SNOMED CT US Edition",
        version="2026-09-01",
        date="2026-09-01",
        file_name="Snomed.zip",
        download_url="https://download.nlm.nih.gov/umls/kss/Snomed.zip",
    )
    return Release(**{**base, **kw})


TYPES = {
    "releaseTypes": [
        {
            "product": "RxNorm",
            "releaseType": "RxNorm Full Monthly Release",
            "endpoint": "https://uts-ws.nlm.nih.gov/releases?releaseType=rxnorm-full-monthly-release",
        },
        {
            "product": "UMLS",
            "releaseType": "UMLS Full Release",
            "endpoint": "https://uts-ws.nlm.nih.gov/releases?releaseType=umls-full-release",
        },
    ]
}


def _listing_client() -> httpx.Client:
    def handler(request: httpx.Request) -> httpx.Response:
        stream = request.url.params.get("releaseType")
        if stream is None:
            return httpx.Response(200, json=TYPES)
        if stream == "umls-full-release":
            return httpx.Response(
                200,
                json=[
                    {
                        "fileName": "umls-2026AA-full.zip",
                        "releaseVersion": "2026AA",
                        "releaseDate": "2026-05-03",
                        "downloadUrl": "https://download.nlm.nih.gov/umls/kss/2026AA/umls-2026AA-full.zip",
                    }
                ],
            )
        return httpx.Response(
            200,
            json=[
                {
                    "fileName": "RxNorm_full_10052026.zip",
                    "releaseDate": "2026-10-05",
                    "downloadUrl": "https://download.nlm.nih.gov/rxnorm/RxNorm_full_10052026.zip",
                }
            ],
        )

    return httpx.Client(transport=httpx.MockTransport(handler))


def test_scrub_removes_key_param_and_literal():
    msg = f"GET https://x/download?url=f.zip&apiKey={KEY} failed; key {KEY}"
    out = scrub(msg, KEY)
    assert KEY not in out
    assert "apiKey=***" in out
    assert scrub("apiKey=abc") == "apiKey=***"


def test_release_types_reads_the_slug_from_each_endpoint():
    with _listing_client() as c:
        types = release_types(c)
    assert [t["stream"] for t in types] == [
        "rxnorm-full-monthly-release",
        "umls-full-release",
    ]
    assert types[1]["name"] == "UMLS Full Release"


def test_current_lists_every_stream_and_falls_back_to_the_date_for_version():
    with _listing_client() as c:
        rels = current(c)
    by_stream = {r.stream: r for r in rels}
    assert by_stream["umls-full-release"].version == "2026AA"
    assert by_stream["rxnorm-full-monthly-release"].version == "2026-10-05"
    assert by_stream["rxnorm-full-monthly-release"].product == "RxNorm"


def test_current_filters_to_named_streams():
    with _listing_client() as c:
        rels = current(c, ["umls-full-release"])
    assert [r.file_name for r in rels] == ["umls-2026AA-full.zip"]


def test_fetch_writes_the_zip_and_a_sidecar_then_reports_present(tmp_path):
    body = _zip_bytes()
    seen = []

    def handler(request: httpx.Request) -> httpx.Response:
        seen.append(request)
        return httpx.Response(200, content=body)

    rel = _rel()
    with httpx.Client(transport=httpx.MockTransport(handler)) as c:
        path, status = fetch(c, rel, tmp_path, KEY)
        again, status2 = fetch(c, rel, tmp_path, KEY)

    assert status == "downloaded" and status2 == "present"
    assert path == target(tmp_path, rel) == again
    assert path.read_bytes() == body
    meta = json.loads(path.with_name(path.name + ".json").read_text())
    assert meta["bytes"] == len(body) and len(meta["sha256"]) == 64
    assert meta["download_url"] == rel.download_url  # the file URL, never the key
    assert KEY not in json.dumps(meta)
    assert len(seen) == 1
    assert seen[0].url.params["url"] == rel.download_url


def test_fetch_resumes_a_partial_download_with_a_range_request(tmp_path):
    body = _zip_bytes()
    rel = _rel()
    path = target(tmp_path, rel)
    path.parent.mkdir(parents=True)
    part = path.with_name(path.name + ".part")
    part.write_bytes(body[:10])

    def handler(request: httpx.Request) -> httpx.Response:
        assert request.headers["Range"] == "bytes=10-"
        return httpx.Response(206, content=body[10:])

    with httpx.Client(transport=httpx.MockTransport(handler)) as c:
        out, status = fetch(c, rel, tmp_path, KEY)
    assert status == "downloaded" and out.read_bytes() == body


def test_fetch_restarts_when_the_server_ignores_the_range(tmp_path):
    body = _zip_bytes()
    rel = _rel()
    path = target(tmp_path, rel)
    path.parent.mkdir(parents=True)
    path.with_name(path.name + ".part").write_bytes(b"stale bytes")

    with httpx.Client(
        transport=httpx.MockTransport(lambda r: httpx.Response(200, content=body))
    ) as c:
        out, _ = fetch(c, rel, tmp_path, KEY)
    assert out.read_bytes() == body


def test_fetch_rejects_a_non_zip_and_never_echoes_the_key(tmp_path):
    html = f"<html>login required for apiKey={KEY}</html>".encode()
    rel = _rel()
    with httpx.Client(
        transport=httpx.MockTransport(lambda r: httpx.Response(200, content=html))
    ) as c:
        with pytest.raises(RuntimeError) as err:
            fetch(c, rel, tmp_path, KEY)
    assert "not a zip" in str(err.value) and KEY not in str(err.value)
    assert not any(tmp_path.rglob("*.part"))


def test_fetch_http_error_is_scrubbed(tmp_path):
    rel = _rel()
    with httpx.Client(
        transport=httpx.MockTransport(lambda r: httpx.Response(401, text="no"))
    ) as c:
        with pytest.raises(RuntimeError) as err:
            fetch(c, rel, tmp_path, KEY)
    assert "401" in str(err.value) and KEY not in str(err.value)


def test_fetch_keeps_non_zip_files_as_is(tmp_path):
    rel = _rel(
        file_name="notes.txt", download_url="https://download.nlm.nih.gov/notes.txt"
    )
    with httpx.Client(
        transport=httpx.MockTransport(lambda r: httpx.Response(200, content=b"plain"))
    ) as c:
        path, status = fetch(c, rel, tmp_path, KEY)
    assert status == "downloaded" and path.read_bytes() == b"plain"