stack/tests/aco/test_load_cms.py

"""aco.load.cms — load data.cms.gov CSVs and derive the CCN ↔ NPI ↔ TIN crosswalk."""

from __future__ import annotations

import csv
from pathlib import Path

import duckdb
import pytest

from aco.load.cms import coverage, derive_crosswalk, load_cms
from cms import datacms
from tests.cms.test_datacms import _catalog

HOSP = [
    (
        "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",
        "SUBGROUP - GENERAL",
        "CAH OR HOSPITAL CCN",
    ),
    (
        "O1",
        "TN",
        "00-09",
        "PART A PROVIDER - HOSPITAL",
        "1467408781",
        "N",
        "440058",
        "5193632180",
        "SOUTHERN TENNESSEE MEDICAL CENTER LLC",
        "HIGHPOINT",
        "11/09/1998",
        "DE",
        "LLC",
        "",
        "P",
        "185 HOSPITAL RD",
        "",
        "WINCHESTER",
        "TN",
        "373982404",
        "Y",
        "440058",
    ),
    (
        "O2",
        "TX",
        "00-09",
        "PART A PROVIDER - HOSPITAL",
        "1000000002",
        "Y",
        "450001",
        "5000000002",
        "TEXAS GENERAL",
        "",
        "",
        "TX",
        "CORP",
        "",
        "N",
        "1 MAIN",
        "",
        "AUSTIN",
        "TX",
        "78701",
        "Y",
        "450001",
    ),
]
SNF = [
    (
        "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",
        "NURSING HOME PROVIDER NAME",
        "AFFILIATION ENTITY NAME",
        "AFFILIATION ENTITY ID",
        "ADDRESS LINE 1",
        "ADDRESS LINE 2",
        "CITY",
        "STATE",
        "ZIP CODE",
    ),
    (
        "S1",
        "TN",
        "00-10",
        "PART A PROVIDER - SNF",
        "1000000003",
        "N",
        "445001",
        "5000000003",
        "WINCHESTER SNF",
        "",
        "",
        "TN",
        "LLC",
        "",
        "P",
        "WINCHESTER SNF",
        "",
        "",
        "2 ELM",
        "",
        "WINCHESTER",
        "TN",
        "37398",
    ),
]
POS = [
    (
        "PRVDR_CTGRY_SBTYP_CD",
        "PRVDR_CTGRY_CD",
        "FAC_NAME",
        "PRVDR_NUM",
        "STATE_CD",
        "PGM_TRMNTN_CD",
        "TRMNTN_EXPRTN_DT",
        "CRTFCTN_DT",
        "ZIP_CD",
    ),
    (
        "01",
        "01",
        "SOUTHERN TENNESSEE MEDICAL CENTER",
        "440058",
        "TN",
        "00",
        "",
        "19900101",
        "37398",
    ),
    (
        "01",
        "01",
        "CLOSED HOSPITAL",
        "440999",
        "TN",
        "01",
        "20200101",
        "19800101",
        "37000",
    ),
]
COST = [
    (
        "rpt_rec_num",
        "Provider CCN",
        "Hospital Name",
        "CCN Facility Type",
        "Type of Control",
    ),
    ("1", "440058", "SOUTHERN TENNESSEE MEDICAL CENTER", "STACH", "Proprietary"),
]


def _write(dir_: Path, name: str, rows: list[tuple]) -> Path:
    p = dir_ / name
    with p.open("w", newline="") as fh:
        csv.writer(fh).writerows(rows)
    return p


@pytest.fixture
def fixtures(tmp_path: Path) -> dict[str, Path]:
    d = tmp_path / "raw"
    d.mkdir()
    files = {
        "hospital_enrollments": _write(d, "Hospital_Enrollments_2026.09.29.csv", HOSP),
        "skilled_nursing_facility_enrollments": _write(
            d, "SNF_Enrollments_2026.09.29.csv", SNF
        ),
        "provider_services_file_quality_improvement_evaluation_system": _write(
            d, "Hospital_and_other.DATA.Q3_2026.csv", POS
        ),
        "hospital_provider_cost_report": _write(d, "CostReport_2021_Final.csv", COST),
    }
    return files


def _catalog_all() -> list[dict]:
    cat = _catalog()
    for table, title in datacms.INSTITUTIONAL.items():
        if title not in {datacms.family_title(d["title"]) for d in cat}:
            cat.append(
                {
                    "title": f"{title} : 2026-10-01",
                    "modified": "2026-10-19",
                    "identifier": "https://data.cms.gov/data-api/v1/dataset/00000000-0000-0000-0000-000000000000/data-viewer",
                    "distribution": [
                        {
                            "mediaType": "text/csv",
                            "downloadURL": f"https://x/{table}.csv",
                        }
                    ],
                }
            )
    return cat


def test_load_and_derive(tmp_path: Path, fixtures: dict[str, Path]) -> None:
    db = tmp_path / "aco.duckdb"
    con = duckdb.connect(str(db))
    con.execute("create schema bcda")
    con.execute(
        "create table bcda.explanation_of_benefit as select * from (values "
        "('INPATIENT', '1467408781', NULL, NULL, NULL, '2026-05-01'), "
        "('INPATIENT', '1000000002', NULL, NULL, NULL, '2026-05-01'), "
        "('INPATIENT', NULL, NULL, NULL, NULL, '2026-05-01'), "
        "('CARRIER', NULL, '1467408781', '621234567', NULL, '2026-06-01'), "
        "('CARRIER', NULL, '1467408781', '621234567', NULL, '2026-07-01'), "
        "('CARRIER', NULL, '1999999999', '000000001', NULL, '2026-07-01'), "
        "('OUTPATIENT', '1467408781', NULL, NULL, '777777777', '2026-07-01')"
        ") t(eob_type, fac_prvdr_npi_num, rndrg_prvdr_npi_num, clm_rndrg_prvdr_tax_num, clm_blg_prvdr_tax_num, last_updated)"
    )
    con.close()

    def fetch(ds, dest_dir, refresh=False):  # noqa: ANN001, ANN202
        return fixtures[ds.table]

    stats = load_cms(
        list(fixtures),
        database=str(db),
        raw_dir=tmp_path / "raw",
        datasets=_catalog_all(),
        fetch=fetch,
    )
    assert (
        stats["cms.hospital_enrollments"] == 2
        and stats["cms.skilled_nursing_facility_enrollments"] == 1
    )
    assert stats["cms.institutional_provider"] == 3
    assert (
        stats["cms.ccn_crosswalk"] == 4
    )  # 3 PECOS rows + the POS-only closed hospital
    assert (
        stats["cms.provider_tin_observed"] == 3
    )  # two rendering pairs + one facility pair
    assert stats["cms.ccn_tin"] == 4

    con = duckdb.connect(str(db), read_only=True)
    cols = {r[0] for r in con.execute("describe cms.hospital_enrollments").fetchall()}
    assert {
        "enrollment_id",
        "ccn",
        "npi",
        "subgroup_general",
        "cah_or_hospital_ccn",
        "_snapshot_date",
        "_source_url",
        "_loaded_at",
    } <= cols
    assert con.execute(
        "select distinct _snapshot_date from cms.hospital_enrollments"
    ).fetchall() == [("2026-09-29",)]
    assert con.execute(
        "select distinct proprietary_nonprofit from cms.institutional_provider order by 1"
    ).fetchall() == [("N",), ("P",)]
    xw = {
        r[0]: r
        for r in con.execute(
            "select ccn, npi, sources, pos_fac_name, pos_termination_cd, cost_type_of_control from cms.ccn_crosswalk"
        ).fetchall()
    }
    assert xw["440058"][1:] == (
        "1467408781",
        "hospital_enrollments",
        "SOUTHERN TENNESSEE MEDICAL CENTER",
        "00",
        "Proprietary",
    )
    assert (
        xw["440999"][1] is None
        and xw["440999"][2] == "provider_of_services"
        and xw["440999"][4] == "01"
    )
    assert xw["445001"][2] == "skilled_nursing_facility_enrollments"
    tin = {
        r[0]: r[1:]
        for r in con.execute(
            "select ccn, tin, tin_source, observed_claims from cms.ccn_tin"
        ).fetchall()
    }
    # the hospital's NPI rendered two carrier lines under one TIN and billed one outpatient claim under another:
    # arg_max by claims picks the two-claim TIN; both are visible in provider_tin_observed
    assert tin["440058"] == ("621234567", "observed_claims", 3)
    assert tin["450001"] == (None, None, None) and tin["440999"] == (None, None, None)
    assert (
        con.execute(
            "select count(*) from cms.ingest_log where module = 'cms'"
        ).fetchone()[0]
        == len(fixtures) + 4
    )
    cov = coverage(con)
    assert cov["available"] and cov["facility_npis"] == 2 and cov["claims"] == 3
    assert (
        cov["npis_with_ccn"] == 2
        and cov["npis_with_tin"] == 1
        and cov["claims_with_tin"] == 2
    )
    assert (
        cov["institutional_claims_without_facility_npi"] == 1
        and cov["ccn_universe"] == 4
    )
    con.close()

    # a participant list wins over observed TINs
    con = duckdb.connect(str(db))
    con.execute(
        "create table cms.aco_participant_tin as select '450001' ccn, NULL::varchar npi, '741111111' tin, 'participant list PY2026' as \"source\""
    )
    derive_crosswalk(con)
    assert con.execute(
        "select tin, tin_source from cms.ccn_tin where ccn = '450001'"
    ).fetchone() == ("741111111", "participant_list")
    con.close()


def test_derive_without_anything(tmp_path: Path) -> None:
    con = duckdb.connect(str(tmp_path / "x.duckdb"))
    con.execute("create schema cms")
    stats = derive_crosswalk(con)
    assert stats == {
        "cms.institutional_provider": 0,
        "cms.ccn_crosswalk": 0,
        "cms.provider_tin_observed": 0,
        "cms.ccn_tin": 0,
    }
    assert coverage(con) == {"available": False}
    con.close()


def test_unknown_table_and_cli(tmp_path: Path, monkeypatch) -> None:
    with pytest.raises(KeyError):
        load_cms(["nope"], database=str(tmp_path / "x.duckdb"), datasets=[])
    from typer.testing import CliRunner

    from cli import app

    monkeypatch.setattr(
        "aco.load.cms.load_cms",
        lambda *a, **k: {"cms.hospital_enrollments": 2, "cms.ccn_tin": 2},
    )
    monkeypatch.setattr(
        "aco.load.cms.coverage",
        lambda con: {
            "available": True,
            "facility_npis": 2,
            "npis_with_ccn": 2,
            "claims_with_ccn": 3,
            "claims": 3,
            "npis_with_tin": 1,
            "claims_with_tin": 2,
            "institutional_claims_without_facility_npi": 1,
            "ccn_universe": 4,
            "observed_tins": 3,
        },
    )
    db = tmp_path / "aco.duckdb"
    duckdb.connect(str(db)).close()
    result = CliRunner().invoke(
        app, ["load", "cms", "--database", str(db), "-t", "hospital_enrollments"]
    )
    assert result.exit_code == 0, result.output
    assert (
        "cms.ccn_tin: 2 rows" in result.output and "2 resolve to a CCN" in result.output
    )