stack/tests/pfs/test_extract.py

"""pfs.extract — deterministic parse + injected classifier → rows and review queue."""

from __future__ import annotations

import sqlite3

import duckdb
import pytest

from pfs.descriptors import DescriptorRun, Para, descriptor_runs
from pfs.extract import (
    Extraction,
    _cpt_elements,
    extract_code,
    extract_run,
    extract_text,
)

STEM = Para(
    "JJ6AM5HJ",
    1163,
    97864,
    "HCPCS code G0556 ( Advanced primary care management services provided by clinical staff and directed by a "
    "physician or other qualified health care professional who is responsible for all primary care and serves as "
    "the continuing focal point for all needed health care services, per calendar month, with the following elements, as appropriate:",
)
ELS = (
    Para("JJ6AM5HJ", 1164, 97864, "Consent;"),
    Para(
        "JJ6AM5HJ",
        1168,
        97864,
        "Provide 24/7 access for urgent needs to care team/practitioner;",
    ),
    Para(
        "JJ6AM5HJ",
        1170,
        97864,
        "Deliver care in alternative ways to traditional office visits to best meet the patient's needs, such as home visits and/or expanded hours;",
    ),
    Para(
        "JJ6AM5HJ",
        1185,
        97865,
        "Be assessed through performance measurement of primary care quality, total cost of care, and meaningful use of Certified EHR Technology).",
    ),
    Para(
        "JJ6AM5HJ", 1199, 97865, "Something the vocabulary does not know about at all;"
    ),
)
RUN = DescriptorRun("G0556", "JJ6AM5HJ", 2025, STEM, ELS)


def _vals(x, type_):
    return sorted(r.value for r in x.rows if r.type == type_)


class TestDeterministic:
    def test_rows_anchor_to_the_paragraph_that_contains_them(self):
        x = extract_run(RUN)
        assert isinstance(x, Extraction) and x.code == "G0556"
        by_value = {r.value: r for r in x.rows}
        assert by_value["consent"].p_id == 1164
        assert by_value["24-7-access"].p_id == 1168
        assert by_value["performance-measurement"].p_id == 1185
        assert by_value["clinical-staff-directed"].p_id == 1163  # stem-level
        assert (
            by_value["calendar-month"].source == "fr"
            and by_value["calendar-month"].year == 2025
        )

    def test_unknown_paragraph_goes_to_review_without_classifier(self):
        x = extract_run(RUN)
        assert [r.p_id for r in x.reviews] == [1199]
        assert x.reviews[0].proposed_value == ""

    def test_home_setting_from_element_line(self):
        assert "home" in _vals(extract_run(RUN), "setting")


class TestClassifier:
    def test_classifier_places_unknown_line(self):
        seen = []

        def classify(text, choices):
            seen.append((text, tuple(choices)))
            return "community-coordination" if "vocabulary" in text else None

        x = extract_run(RUN, classify=classify)
        assert x.reviews == ()
        row = next(r for r in x.rows if r.value == "community-coordination")
        assert row.type == "activity" and row.p_id == 1199
        assert seen and "consent" in seen[0][1] and "significant-risk" in seen[0][1]

    def test_classifier_none_means_review_with_no_proposal(self):
        x = extract_run(RUN, classify=lambda t, c: None)
        assert len(x.reviews) == 1 and x.reviews[0].proposed_value == ""

    def test_classifier_not_called_for_deterministic_lines(self):
        calls = []
        extract_run(RUN, classify=lambda t, c: calls.append(t) or None)
        assert calls == ["Something the vocabulary does not know about at all;"]


class TestAnchorIntegrity:
    def test_no_row_anchored_to_a_foreign_code_paragraph(self):
        # C1 live-corpus incident: an enumeration/cross-reference line for
        # a *different* code ("( 9) 99457 and 99458 (codes for remote
        # physiologic monitoring, each additional 20 minutes).") must
        # never contribute an ElementRow to a 99439 run. The fix lives in
        # `descriptor_runs`/`is_element_paragraph` (it stops the run
        # before absorbing such a line) — this test drives the real
        # `descriptor_runs` to build the run, then checks `extract_run`'s
        # output end to end: no row may be anchored to a paragraph whose
        # text contains a different 5-char code and none of 99439.
        con = sqlite3.connect(":memory:")
        con.row_factory = sqlite3.Row
        con.executescript(
            "CREATE TABLE items (key TEXT PRIMARY KEY, title TEXT, date_published TEXT);"
            "CREATE TABLE fr_anchors (item_key TEXT, p_id INTEGER, page INTEGER, ordinal INTEGER, text TEXT);"
        )
        con.execute(
            "INSERT INTO items VALUES ('XFGGRBDH', "
            "'Medicare and Medicaid Programs; CY 2025 Payment Policies', '2024-07-31')"
        )
        rows = [
            (
                832,
                61652,
                "CPT code 99439 (Chronic care management services, each additional 20 "
                "minutes of clinical staff time directed by a physician or other "
                "qualified health care professional, per calendar month, with the "
                "following required elements:",
            ),
            (833, 61652, "Consent;"),
            (
                834,
                61652,
                "( 9) 99457 and 99458 (codes for remote physiologic monitoring, "
                "each additional 20 minutes).",
            ),
        ]
        con.executemany(
            "INSERT INTO fr_anchors VALUES (?,?,?,?,?)",
            [("XFGGRBDH", p, pg, p, t) for p, pg, t in rows],
        )

        class _S:
            def _con(self):
                return con

        try:
            run = descriptor_runs(_S(), "99439")[0]
            # the enumeration line for 99457/99458 never joins the run
            assert [p.p_id for p in run.elements] == [833]
            x = extract_run(run)
            for r in x.rows:
                assert "99457" not in r.text and "99458" not in r.text
        finally:
            con.close()


def _sqlite_store_con():
    con = sqlite3.connect(":memory:")
    con.row_factory = sqlite3.Row
    con.executescript(
        "CREATE TABLE items (key TEXT PRIMARY KEY, title TEXT, date_published TEXT);"
        "CREATE TABLE fr_anchors (item_key TEXT, p_id INTEGER, page INTEGER, ordinal INTEGER, text TEXT);"
    )
    return con


class _Store:
    def __init__(self, con):
        self._c = con

    def _con(self):
        return self._c


def _duckdb_con():
    c = duckdb.connect(":memory:")
    c.execute("CREATE SCHEMA pfs")
    c.execute(
        "CREATE TABLE pfs.rvu (hcpcs VARCHAR, mod VARCHAR, description VARCHAR, "
        "status_code VARCHAR, non_fac_total DOUBLE, year INTEGER)"
    )
    c.execute("CREATE SCHEMA terminology")
    c.execute(
        "CREATE TABLE terminology.hcpcs_level_2 (hcpcs VARCHAR, long_description VARCHAR, "
        "seqnum VARCHAR, recid VARCHAR)"
    )
    c.execute(
        "CREATE TABLE pfs.cpt_code (edition_year INTEGER, item_key VARCHAR, code VARCHAR, "
        "sec_id VARCHAR, category VARCHAR, descriptor VARCHAR, stem VARCHAR, "
        "elements VARCHAR[], tail VARCHAR, parent VARCHAR, addon BOOLEAN, resequenced BOOLEAN, "
        "new BOOLEAN, revised BOOLEAN, telemedicine BOOLEAN, mod51_exempt BOOLEAN, "
        "audio_only BOOLEAN, fda_pending BOOLEAN, pla BOOLEAN)"
    )
    return c


def _duckdb_con_no_cpt():
    """A replica shape that hasn't been through `cpt-ingest` yet: `pfs.rvu`
    and `terminology.hcpcs_level_2` exist, but no `pfs.cpt_*` table does —
    the read-only path deliberately never calls `ensure_tables` (I4)."""
    c = duckdb.connect(":memory:")
    c.execute("CREATE SCHEMA pfs")
    c.execute(
        "CREATE TABLE pfs.rvu (hcpcs VARCHAR, mod VARCHAR, description VARCHAR, "
        "status_code VARCHAR, non_fac_total DOUBLE, year INTEGER)"
    )
    c.execute("CREATE SCHEMA terminology")
    c.execute(
        "CREATE TABLE terminology.hcpcs_level_2 (hcpcs VARCHAR, long_description VARCHAR, "
        "seqnum VARCHAR, recid VARCHAR)"
    )
    return c


def _insert_cpt_code(con, year, item_key, code, *, stem, elements, tail):
    con.execute(
        "INSERT INTO pfs.cpt_code VALUES (?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?)",
        [
            year,
            item_key,
            code,
            "sec1",
            "I",
            f"{stem}; {tail}",
            stem,
            list(elements),
            tail,
            "",
            False,
            False,
            False,
            False,
            False,
            False,
            False,
            False,
            False,
        ],
    )


class TestExtractCodeCpt:
    @pytest.fixture
    def store(self):
        con = _sqlite_store_con()
        yield _Store(con)
        con.close()

    @pytest.fixture
    def con(self):
        c = _duckdb_con()
        yield c
        c.close()

    def test_cpt_row_yields_source_cpt_rows_and_a_review_row(self, store, con):
        _insert_cpt_code(
            con,
            2024,
            "GQGTPGYV",
            "99490",
            stem="Chronic care management services",
            elements=[
                "Consent;",
                "Something the vocabulary does not know about at all;",
            ],
            tail="first 20 minutes, per calendar month.",
        )
        x = extract_code(store, con, "99490")
        cpt_rows = [r for r in x.rows if r.source == "cpt"]
        assert cpt_rows
        assert all(r.item_key == "GQGTPGYV" and r.year == 2024 for r in cpt_rows)
        assert any(r.value == "consent" for r in cpt_rows)
        stem_row = next(r for r in cpt_rows if r.value == "calendar-month")
        assert stem_row.text == "Chronic care management services"
        assert stem_row.p_id == 0 and stem_row.page == 0
        assert [(r.text, r.item_key, r.p_id) for r in x.reviews] == [
            (
                "Something the vocabulary does not know about at all",
                "GQGTPGYV",
                0,
            )
        ]

    def test_no_cpt_code_row_yields_no_cpt_rows(self, store, con):
        x = extract_code(store, con, "99490")
        assert not any(r.source == "cpt" for r in x.rows)
        assert x.reviews == ()

    def test_replica_with_no_cpt_tables_at_all_does_not_raise(self, store):
        # #685/#686 review finding: pfs.cpt_code may not exist yet (an
        # older replica, or one that hasn't been through cpt-ingest) —
        # the read-only `elements --dry-run` path never calls
        # ensure_tables, so this must not raise CatalogException.
        con = _duckdb_con_no_cpt()
        try:
            x = extract_code(store, con, "99490")
        finally:
            con.close()
        assert not any(r.source == "cpt" for r in x.rows)
        assert x.reviews == ()

    def test_fr_wins_over_cpt_on_duplicate(self, con):
        store_con = _sqlite_store_con()
        try:
            store_con.execute(
                "INSERT INTO items VALUES ('XFGGRBDH', "
                "'Medicare and Medicaid Programs; CY 2025 Payment Policies', '2024-07-31')"
            )
            rows = [
                (
                    900,
                    60000,
                    "CPT code 99490 (Chronic care management services, with the "
                    "following required elements:",
                ),
                (901, 60000, "Consent;"),
            ]
            store_con.executemany(
                "INSERT INTO fr_anchors VALUES (?,?,?,?,?)",
                [("XFGGRBDH", p, pg, p, t) for p, pg, t in rows],
            )
            _insert_cpt_code(
                con,
                2024,
                "GQGTPGYV",
                "99490",
                stem="Chronic care management services",
                elements=["Consent;"],
                tail="with the following required elements.",
            )
            x = extract_code(_Store(store_con), con, "99490")
            consent_rows = [r for r in x.rows if r.value == "consent"]
            assert len(consent_rows) == 1
            assert consent_rows[0].source == "fr" and consent_rows[0].p_id == 901
        finally:
            store_con.close()

    def test_cpt_wins_over_hcpcs_on_duplicate(self, store, con):
        con.execute(
            "INSERT INTO terminology.hcpcs_level_2 VALUES (?,?,?,?)",
            ["99490", "… per calendar month, consent …", "1", "1"],
        )
        _insert_cpt_code(
            con,
            2024,
            "GQGTPGYV",
            "99490",
            stem="Chronic care management services",
            elements=["Consent;"],
            tail="per calendar month.",
        )
        x = extract_code(store, con, "99490")
        consent_rows = [r for r in x.rows if r.value == "consent"]
        assert len(consent_rows) == 1
        assert consent_rows[0].source == "cpt"

    def test_classifier_places_an_unmatched_cpt_element_line(self, store, con):
        # Mirrors TestClassifier's FR-side coverage, but for a required-
        # elements list item nothing deterministic can place — the
        # classify(text, _CHOICES) branch inside _cpt_elements.
        _insert_cpt_code(
            con,
            2024,
            "GQGTPGYV",
            "99490",
            stem="Chronic care management services",
            elements=[
                "Something the vocabulary does not know about at all;",
            ],
            tail="per calendar month.",
        )
        classify = lambda text, choices: (  # noqa: E731
            "community-coordination" if "vocabulary" in text else None
        )
        x = extract_code(store, con, "99490", classify=classify)
        cpt_rows = [r for r in x.rows if r.source == "cpt"]
        row = next(r for r in cpt_rows if r.value == "community-coordination")
        assert row.type == "activity" and row.item_key == "GQGTPGYV"
        assert not x.reviews

    def test_blank_cpt_element_after_stripping_punctuation_is_skipped(self, store, con):
        # An element list item that is only trailing punctuation (a
        # stray ";" from the codebook's own formatting) must not become
        # a row or a review once stripped down to "".
        _insert_cpt_code(
            con,
            2024,
            "GQGTPGYV",
            "99490",
            stem="Chronic care management services",
            elements=["Consent;", ";", "   ;  "],
            tail="per calendar month.",
        )
        x = extract_code(store, con, "99490")
        cpt_rows = [r for r in x.rows if r.source == "cpt"]
        assert {r.value for r in cpt_rows} >= {"consent", "calendar-month"}
        assert not x.reviews


class TestCptElementsErrors:
    class _RaisingCon:
        def execute(self, *args, **kwargs):
            raise RuntimeError("disk I/O error")

    def test_non_missing_table_error_propagates(self):
        # is_missing_table_error only swallows a DuckDB "Catalog ...
        # does not exist" error — any other failure querying
        # pfs.cpt_code (a real I/O error, a corrupt replica) must
        # propagate rather than silently degrade to no CPT rows.
        with pytest.raises(RuntimeError, match="disk I/O error"):
            _cpt_elements(self._RaisingCon(), "99490")


class TestExtractCodeRvu:
    @pytest.fixture
    def store(self):
        con = _sqlite_store_con()
        yield _Store(con)
        con.close()

    @pytest.fixture
    def con(self):
        c = _duckdb_con()
        yield c
        c.close()

    def test_rvu_description_rows_are_merged_per_year(self, store, con):
        # extract_code's final loop over rvu_descriptions() — a source
        # not exercised by the hcpcs-long-description tests above. Each
        # year's description names a distinct element so neither is
        # dropped by the (type, value, detail) dedupe merged.setdefault
        # already applies within a single source.
        con.execute(
            "INSERT INTO pfs.rvu VALUES (?,?,?,?,?,?)",
            ["99490", "", "… per calendar month, consent …", "A", 1.0, 2023],
        )
        con.execute(
            "INSERT INTO pfs.rvu VALUES (?,?,?,?,?,?)",
            ["99490", "", "… provide 24/7 access for urgent needs …", "A", 1.0, 2024],
        )
        x = extract_code(store, con, "99490")
        rvu_rows = [r for r in x.rows if r.source == "rvu"]
        by_value = {r.value: r for r in rvu_rows}
        assert by_value["consent"].year == 2023
        assert by_value["24-7-access"].year == 2024


class TestText:
    def test_hcpcs_long_description(self):
        x = extract_text(
            "G0556",
            "… per calendar month, with the following elements: consent; 24/7 access …",
            year=2025,
            source="hcpcs",
        )
        assert {r.value for r in x.rows} >= {"calendar-month", "consent", "24-7-access"}
        assert all(
            r.item_key == "" and r.p_id == 0 and r.source == "hcpcs" for r in x.rows
        )