stack/tests/cli/test_nlm_cli.py

"""stack nlm — releases / fetch / vsac."""

from __future__ import annotations

import gzip
import json

import pytest
from typer.testing import CliRunner

import cli.nlm as nlm_cli
import nlm.releases as releases_mod
import nlm.vsac as vsac_mod
from cli import app
from nlm.releases import Release

runner = CliRunner()

REL = Release(
    product="UMLS",
    stream="umls-full-release",
    name="UMLS Full Release",
    version="2026AA",
    date="2026-05-03",
    file_name="umls-2026AA-full.zip",
    download_url="https://download.nlm.nih.gov/umls/kss/2026AA/umls-2026AA-full.zip",
)


@pytest.fixture
def key(monkeypatch):
    monkeypatch.setenv("UMLS_API_KEY", "k-123")


def test_missing_key_exits_with_a_hint(monkeypatch, tmp_path):
    monkeypatch.delenv("UMLS_API_KEY", raising=False)
    monkeypatch.setattr("conf.secret", lambda *a: "")
    r = runner.invoke(app, ["nlm", "fetch", "--dest", str(tmp_path)])
    assert r.exit_code == 2 and "UMLS_API_KEY" in r.output


def test_releases_lists_current_files(monkeypatch):
    monkeypatch.setattr(releases_mod, "current", lambda client, streams: [REL])
    r = runner.invoke(app, ["nlm", "releases"])
    assert r.exit_code == 0 and "umls-2026AA-full.zip" in r.output


def test_fetch_reports_each_file_and_fails_when_one_did(monkeypatch, tmp_path, key):
    other = Release(
        **{
            **REL.__dict__,
            "stream": "rxnorm-full-monthly-release",
            "file_name": "rx.zip",
        }
    )
    monkeypatch.setattr(releases_mod, "current", lambda client, streams: [REL, other])

    def fake_fetch(client, rel, root, api_key):
        assert api_key == "k-123"
        if rel.file_name == "rx.zip":
            raise RuntimeError("rx.zip: 401 Unauthorized")
        p = root / rel.stream / rel.version / rel.file_name
        p.parent.mkdir(parents=True)
        p.write_bytes(b"x" * 2_000_000)
        return p, "downloaded"

    monkeypatch.setattr(releases_mod, "fetch", fake_fetch)
    r = runner.invoke(app, ["nlm", "fetch", "--dest", str(tmp_path)])
    assert r.exit_code == 1
    assert "downloaded" in r.output and "2.0 MB" in r.output
    assert "failed     rx.zip: 401" in r.output


def test_fetch_defaults_to_the_configured_dir(monkeypatch, tmp_path, key):
    monkeypatch.setattr(releases_mod, "current", lambda client, streams: [])
    monkeypatch.setattr("conf.path", lambda k: tmp_path / k)
    assert nlm_cli._root(None) == tmp_path / "nlm.dir"
    r = runner.invoke(app, ["nlm", "fetch"])
    assert r.exit_code == 0


def _stub_vsac(monkeypatch, root, ids):
    monkeypatch.setattr(vsac_mod, "make_client", lambda k: _Dummy())

    def save(client, r):
        r.mkdir(parents=True, exist_ok=True)
        (r / "index.jsonl").write_text(
            "".join(json.dumps({"id": i}) + "\n" for i in ids) + "\n"
        )
        return len(ids)

    monkeypatch.setattr(vsac_mod, "save_definitions", save)


class _Dummy:
    def __enter__(self):
        return self

    def __exit__(self, *a):
        return False


def test_vsac_definitions_only(monkeypatch, tmp_path, key):
    _stub_vsac(monkeypatch, tmp_path, ["a", "b"])
    monkeypatch.setattr(
        vsac_mod, "hoard", lambda *a, **k: pytest.fail("should not expand")
    )
    r = runner.invoke(
        app, ["nlm", "vsac", "--dest", str(tmp_path), "--definitions-only"]
    )
    assert r.exit_code == 0 and "definitions: 2" in r.output


def test_vsac_expands_with_limit_and_counts_codes(monkeypatch, tmp_path, key):
    _stub_vsac(monkeypatch, tmp_path, [f"id{i}" for i in range(600)])
    seen = {}

    def fake_hoard(factory, root, ids, *, api_key, workers, progress):
        seen.update(ids=ids, workers=workers, api_key=api_key, client=factory())
        exp = root / "expansions"
        exp.mkdir(parents=True)
        with gzip.open(exp / "id0.json.gz", "wt") as f:
            json.dump({"expansion": {"contains": [{"code": "1"}, {"code": "2"}]}}, f)
        for done in range(1, 501):
            progress(done, 500)
        return {"expanded": 500, "present": 0, "failed": 0}

    monkeypatch.setattr(vsac_mod, "hoard", fake_hoard)
    r = runner.invoke(
        app,
        ["nlm", "vsac", "--dest", str(tmp_path), "--limit", "500", "--workers", "2"],
    )
    assert r.exit_code == 0, r.output
    assert (
        len(seen["ids"]) == 500 and seen["workers"] == 2 and seen["api_key"] == "k-123"
    )
    assert isinstance(seen["client"], _Dummy)
    assert "expanded 500/500" in r.output
    assert "expansions: 500 new, 0 present, 0 failed" in r.output
    assert "codes on disk: 2" in r.output