stack/tests/aco/test_express_hcc_suspecting.py

"""Tests for aco.express.hcc_suspecting — HCC suspecting staging functions."""

from __future__ import annotations

import polars as pl

from aco.express.hcc_suspecting import (
    stg_condition,
    stg_lab_result,
    stg_medication,
    stg_observation,
    stg_patient,
    stg_pharmacy_claim,
)

# ── stg_condition ────────────────────────────────────────────────────────────


class TestStgCondition:
    def test_returns_dataframe(self, condition_df) -> None:
        result = stg_condition(condition_df)
        assert isinstance(result, pl.DataFrame)

    def test_selects_expected_columns(self, condition_df) -> None:
        result = stg_condition(condition_df)
        for col in [
            "claim_id",
            "person_id",
            "recorded_date",
            "code_type",
            "code",
            "data_source",
        ]:
            assert col in result.columns

    def test_lowercases_code_type(self) -> None:
        df = pl.DataFrame(
            {
                "condition_id": ["C1"],
                "person_id": ["P1"],
                "patient_id": [None],
                "encounter_id": [None],
                "claim_id": ["CLM1"],
                "recorded_date": [None],
                "onset_date": [None],
                "resolved_date": [None],
                "status": [None],
                "condition_type": ["problem"],
                "source_code_type": [None],
                "source_code": [None],
                "source_description": [None],
                "normalized_code_type": ["ICD-10-CM"],
                "normalized_code": ["E11.9"],
                "normalized_description": [None],
                "condition_rank": [1],
                "present_on_admit_code": [None],
                "present_on_admit_description": [None],
                "data_source": ["test"],
                "file_name": [None],
                "ingest_datetime": [None],
                "payer": ["Medicare"],
            }
        )
        result = stg_condition(df)
        assert result["code_type"][0] == "icd-10-cm"

    def test_renames_columns(self, condition_df) -> None:
        result = stg_condition(condition_df)
        assert "code" in result.columns
        assert "normalized_code" not in result.columns

    def test_row_count_preserved(self, condition_df) -> None:
        result = stg_condition(condition_df)
        assert len(result) == len(condition_df)


# ── stg_lab_result ───────────────────────────────────────────────────────────


class TestStgLabResult:
    def test_passthrough(self) -> None:
        df = pl.DataFrame({"lab_id": ["L1"], "value": [1.5], "data_source": ["test"]})
        result = stg_lab_result(df)
        assert isinstance(result, pl.DataFrame)
        assert len(result) == 1
        assert result.columns == df.columns


# ── stg_medication ───────────────────────────────────────────────────────────


class TestStgMedication:
    def test_selects_expected_columns(self, medication_df) -> None:
        result = stg_medication(medication_df)
        assert isinstance(result, pl.DataFrame)
        for col in [
            "person_id",
            "dispensing_date",
            "source_code",
            "ndc_code",
            "rxnorm_code",
            "data_source",
        ]:
            assert col in result.columns

    def test_drops_extra_columns(self, medication_df) -> None:
        result = stg_medication(medication_df)
        assert "medication_id" not in result.columns
        assert "route" not in result.columns


# ── stg_observation ──────────────────────────────────────────────────────────


class TestStgObservation:
    def test_selects_expected_columns(self, observation_df) -> None:
        result = stg_observation(observation_df)
        assert isinstance(result, pl.DataFrame)
        for col in [
            "person_id",
            "payer",
            "observation_date",
            "result",
            "code_type",
            "code",
            "data_source",
        ]:
            assert col in result.columns

    def test_payer_literal(self, observation_df) -> None:
        result = stg_observation(observation_df)
        assert result["payer"][0] == "clinical source"

    def test_fallback_code_type(self) -> None:
        df = pl.DataFrame(
            {
                "observation_id": ["O1"],
                "person_id": ["P1"],
                "patient_id": [None],
                "encounter_id": [None],
                "panel_id": [None],
                "observation_date": [None],
                "observation_type": [None],
                "source_code_type": ["loinc"],
                "source_code": ["1234-5"],
                "source_description": [None],
                "normalized_code_type": [None],
                "normalized_code": [None],
                "normalized_description": [None],
                "result": ["120"],
                "source_units": [None],
                "normalized_units": [None],
                "source_reference_range_low": [None],
                "source_reference_range_high": [None],
                "normalized_reference_range_low": [None],
                "normalized_reference_range_high": [None],
                "data_source": ["test"],
            }
        )
        result = stg_observation(df)
        assert result["code_type"][0] == "loinc"
        assert result["code"][0] == "1234-5"


# ── stg_patient ──────────────────────────────────────────────────────────────


class TestStgPatient:
    def test_selects_expected_columns(self, patient_df) -> None:
        result = stg_patient(patient_df)
        assert isinstance(result, pl.DataFrame)
        for col in ["person_id", "sex", "birth_date", "death_date"]:
            assert col in result.columns

    def test_drops_extra_columns(self, patient_df) -> None:
        result = stg_patient(patient_df)
        assert "first_name" not in result.columns
        assert "address" not in result.columns

    def test_row_count_preserved(self, patient_df) -> None:
        result = stg_patient(patient_df)
        assert len(result) == len(patient_df)


# ── stg_pharmacy_claim ───────────────────────────────────────────────────────


class TestStgPharmacyClaim:
    def test_selects_expected_columns(self, pharmacy_claim_df) -> None:
        result = stg_pharmacy_claim(pharmacy_claim_df)
        assert isinstance(result, pl.DataFrame)
        for col in [
            "person_id",
            "payer",
            "dispensing_date",
            "ndc_code",
            "paid_date",
            "data_source",
        ]:
            assert col in result.columns

    def test_drops_extra_columns(self, pharmacy_claim_df) -> None:
        result = stg_pharmacy_claim(pharmacy_claim_df)
        assert "claim_id" not in result.columns
        assert "quantity" not in result.columns