stack/tests/cms/test_datacms.py

"""cms.datacms — catalog resolution, column naming, download paths."""

from __future__ import annotations

from pathlib import Path

import httpx
import pytest

from cms import datacms

HOSPITAL_HEADER = [
    "ENROLLMENT ID",
    "ENROLLMENT STATE",
    "PROVIDER TYPE CODE",
    "PROVIDER TYPE TEXT",
    "NPI",
    "MULTIPLE NPI FLAG",
    "CCN",
    "ASSOCIATE ID",
    "ORGANIZATION NAME",
    "DOING BUSINESS AS NAME",
    "INCORPORATION DATE",
    "INCORPORATION STATE",
    "ORGANIZATION TYPE STRUCTURE",
    "ORGANIZATION OTHER TYPE TEXT",
    "PROPRIETARY NONPROFIT",
    "ADDRESS LINE 1",
    "ADDRESS LINE 2",
    "CITY",
    "STATE",
    "ZIP CODE",
    "PRACTICE LOCATION TYPE",
    "LOCATION OTHER TYPE TEXT",
    "SUBGROUP - GENERAL",
    "SUBGROUP - ACUTE CARE",
    "SUBGROUP - ALCOHOL DRUG",
    "SUBGROUP - CHILDREN",
]


def test_column_names_match_generated_models() -> None:
    from cms.table.hospital_enrollments import HospitalEnrollments
    from cms.table.hospital_provider_cost_report import HospitalProviderCostReport
    from cms.table.skilled_nursing_facility_enrollments import (
        SkilledNursingFacilityEnrollments,
    )

    fields = set(HospitalEnrollments.model_fields)
    # the generated model predates a few SUBGROUP columns CMS added; the
    # loader takes its columns from the CSV, so only the stable core must map
    for h in HOSPITAL_HEADER[:22]:
        assert datacms.sanitize_column_name(h) in fields, h
    assert datacms.sanitize_column_name("PROPRIETARY_NONPROFIT") in set(
        SkilledNursingFacilityEnrollments.model_fields
    )
    cost = set(HospitalProviderCostReport.model_fields)
    for h in (
        "rpt_rec_num",
        "Provider CCN",
        "Hospital Name",
        "CCN Facility Type",
        "Type of Control",
    ):
        assert datacms.sanitize_column_name(h) in cost, h
    assert datacms.sanitize_column_name("3-Day Rule") == "_3_day_rule"
    assert datacms.sanitize_column_name("Type") == "type_"


def test_family_title_and_snapshot_date() -> None:
    assert (
        datacms.family_title("Hospital Enrollments : 2026-10-01")
        == "Hospital Enrollments"
    )
    assert datacms.family_title("Hospital Enrollments") == "Hospital Enrollments"
    assert (
        datacms.snapshot_date_from_name("Hospital_Enrollments_2026.09.29.csv")
        == "2026-09-29"
    )
    assert (
        datacms.snapshot_date_from_name("Hospital_and_other.DATA.Q3_2026.csv")
        == "2026-Q3"
    )
    assert datacms.snapshot_date_from_name("CostReport_2021_Final.csv") == "2021"
    assert datacms.snapshot_date_from_name("nothing") == ""


def _catalog() -> list[dict]:
    fam = "f6f6505c-e8b0-4d57-b258-e2b94133aaf2"
    return [
        {
            "title": "Hospital Enrollments : 2026-09-01",
            "modified": "2026-09-03",
            "identifier": f"https://data.cms.gov/data-api/v1/dataset/{fam}/data-viewer",
            "distribution": [
                {
                    "mediaType": "text/csv",
                    "downloadURL": "https://x/old/Hospital_Enrollments_2026.08.29.csv",
                },
                {"accessURL": f"https://data.cms.gov/data-api/v1/dataset/{fam}/data"},
            ],
        },
        {
            "title": "Hospital Enrollments : 2026-10-01",
            "modified": "2026-10-19",
            "identifier": f"https://data.cms.gov/data-api/v1/dataset/{fam}/data-viewer",
            "distribution": [
                {"accessURL": f"https://data.cms.gov/data-api/v1/dataset/{fam}/data"},
                {
                    "mediaType": "text/csv",
                    "downloadURL": "https://x/new/Hospital_Enrollments_2026.09.29.csv",
                },
                {
                    "accessURL": "https://data.cms.gov/data-api/v1/dataset/71b9466f-aaaa-bbbb-cccc-000000000000/data"
                },
            ],
        },
        {
            "title": "Hospital Provider Cost Report : 2021-12-05",
            "modified": "2026-10-05",
            "identifier": "https://data.cms.gov/data-api/v1/dataset/98b7383e-ad67-44d5-a7d8-0a80a92a72c1/data-viewer",
            "distribution": [
                {
                    "accessURL": "https://data.cms.gov/data-api/v1/dataset/98b7383e-ad67-44d5-a7d8-0a80a92a72c1/data"
                }
            ],
        },
    ]


def test_resolve_picks_newest_version_and_family_api() -> None:
    ds = datacms.resolve("hospital_enrollments", _catalog())
    assert ds.version_title.endswith("2026-10-01") and ds.modified == "2026-10-19"
    assert ds.csv_url.endswith("Hospital_Enrollments_2026.09.29.csv")
    assert ds.api_url.endswith("f6f6505c-e8b0-4d57-b258-e2b94133aaf2/data")
    assert (
        ds.filename == "Hospital_Enrollments_2026.09.29.csv"
        and ds.snapshot_date == "2026-09-29"
    )
    cost = datacms.resolve("hospital_provider_cost_report", _catalog())
    assert (
        cost.csv_url == "" and cost.filename == "hospital_provider_cost_report.api.csv"
    )
    with pytest.raises(LookupError):
        datacms.resolve("hospice_enrollments", _catalog())
    with pytest.raises(KeyError):
        datacms.resolve("nope", _catalog())
    assert (
        datacms.resolve("nope", _catalog(), title="Hospital Enrollments").table
        == "nope"
    )


def test_catalog_cache(tmp_path: Path, monkeypatch) -> None:
    calls = []

    class R:
        def json(self):  # noqa: ANN202
            calls.append(1)
            return {"dataset": _catalog()}

    monkeypatch.setattr(datacms, "_get", lambda url, **kw: R())
    cache = tmp_path / "catalog.json"
    assert len(datacms.catalog(cache)) == 3 and cache.exists() and calls == [1]
    assert len(datacms.catalog(cache)) == 3 and calls == [1]  # served from disk
    assert len(datacms.catalog(cache, refresh=True)) == 3 and calls == [1, 1]
    assert len(datacms.catalog(None)) == 3


def test_download_csv_and_api_paging(tmp_path: Path, monkeypatch) -> None:
    ds = datacms.resolve("hospital_enrollments", _catalog())

    class Stream:
        def __init__(self, *a, **k):  # noqa: ANN002, ANN003
            pass

        def __enter__(self):  # noqa: ANN204
            return self

        def __exit__(self, *a):  # noqa: ANN002, ANN204
            return False

        def raise_for_status(self) -> None:
            pass

        def iter_bytes(self, n):  # noqa: ANN001, ANN202
            yield b'"CCN","NPI"\n'
            yield b"440058,1467408781\n"

    monkeypatch.setattr(datacms.httpx, "stream", lambda *a, **k: Stream())
    p = datacms.download(ds, tmp_path)
    assert p.name == ds.filename and p.read_text().count("\n") == 2
    # cached: the stream is not opened again
    monkeypatch.setattr(
        datacms.httpx,
        "stream",
        lambda *a, **k: (_ for _ in ()).throw(AssertionError("re-download")),
    )
    assert datacms.download(ds, tmp_path) == p

    cost = datacms.resolve("hospital_provider_cost_report", _catalog())
    pages = [
        [{"Provider CCN": "010001", "Hospital Name": "A"}] * datacms.PAGE_SIZE,
        [{"Provider CCN": "010002", "Hospital Name": "B"}],
    ]

    class J:
        def __init__(self, rows):  # noqa: ANN001
            self.rows = rows

        def json(self):  # noqa: ANN202
            return self.rows

    seen = []

    def fake_get(url, **kw):  # noqa: ANN001, ANN003, ANN202
        seen.append(kw["params"]["offset"])
        return J(pages.pop(0) if pages else [])

    monkeypatch.setattr(datacms, "_get", fake_get)
    p = datacms.download(cost, tmp_path)
    assert seen == [0, datacms.PAGE_SIZE]
    lines = p.read_text().splitlines()
    assert (
        lines[0] == "Provider CCN,Hospital Name" and len(lines) == datacms.PAGE_SIZE + 2
    )


def test_get_retries(monkeypatch) -> None:
    attempts = []

    def fake(url, **kw):  # noqa: ANN001, ANN003, ANN202
        attempts.append(url)
        if len(attempts) == 1:
            raise httpx.ConnectError("boom")
        if len(attempts) == 2:
            return httpx.Response(503, request=httpx.Request("GET", url))
        return httpx.Response(200, json={"ok": 1}, request=httpx.Request("GET", url))

    monkeypatch.setattr(datacms.httpx, "get", fake)
    monkeypatch.setattr(datacms.time, "sleep", lambda s: None)
    assert datacms._get("https://x").json() == {"ok": 1} and len(attempts) == 3