stack/tests/aco/test_express_pharmacy.py

"""Tests for aco.express.pharmacy — pharmacy pipeline functions.

Tests verify brand/generic identification, cost calculations, and
aggregation logic.
"""

from __future__ import annotations

import polars as pl
import pytest

from aco.express.pharmacy import (
    int_brand_with_generic_available,
    int_claims_current_cost,
    int_generic_cost,
    int_generic_cost_by_ndc,
)

# ── fixtures ─────────────────────────────────────────────────────────────────


@pytest.fixture
def rxnorm_generic_df() -> pl.DataFrame:
    """RxNorm generic availability lookup."""
    return pl.DataFrame(
        {
            "product_rxcui": ["RX001", "RX001", "RX002"],
            "ndc_product_tty": ["SCD", "GPCK", "SBD"],
            "product_startmarketingdate": [
                "2020-01-01",
                "2021-06-15",
                "2023-01-01",
            ],
            "ndc": ["NDC001", "NDC002", "NDC003"],
        }
    )


@pytest.fixture
def pharmacy_claim_df() -> pl.DataFrame:
    """Core pharmacy claims."""
    return pl.DataFrame(
        {
            "claim_id": ["C1", "C2", "C3"],
            "ndc_code": ["NDC001", "NDC001", "NDC003"],
            "paid_amount": [10.0, 15.0, 50.0],
            "quantity": [30.0, 60.0, 30.0],
            "data_source": ["test", "test", "test"],
            "person_id": ["P1", "P1", "P2"],
        }
    )


@pytest.fixture
def ndc_terminology_df() -> pl.DataFrame:
    """NDC terminology lookup."""
    return pl.DataFrame(
        {
            "ndc": ["NDC001", "NDC002", "NDC003"],
            "fda_description": ["Generic Drug A", "Generic Drug A", "Brand Drug B"],
            "rxcui": ["RX001", "RX001", "RX002"],
        }
    )


@pytest.fixture
def rxnorm_brand_generic_df() -> pl.DataFrame:
    """RxNorm brand/generic mapping."""
    return pl.DataFrame(
        {
            "product_rxcui": ["RX001", "RX002"],
            "brand_vs_generic": ["generic", "brand"],
        }
    )


# ── int_brand_with_generic_available ─────────────────────────────────────────


class TestIntBrandWithGenericAvailable:
    def test_filters_scd_gpck(self, rxnorm_generic_df) -> None:
        result = int_brand_with_generic_available(rxnorm_generic_df)
        assert isinstance(result, pl.DataFrame)
        # Should include SCD and GPCK but not SBD
        assert len(result) <= 2
        assert "brand_with_generic_available" in result.columns

    def test_unique_rxcui(self, rxnorm_generic_df) -> None:
        result = int_brand_with_generic_available(rxnorm_generic_df)
        # RX001 appears twice (SCD + GPCK) but should deduplicate to 1
        assert result["brand_with_generic_available"].n_unique() <= 1

    def test_excludes_future_dates(self) -> None:
        df = pl.DataFrame(
            {
                "product_rxcui": ["RX001"],
                "ndc_product_tty": ["SCD"],
                "product_startmarketingdate": ["2099-01-01"],
                "ndc": ["NDC001"],
            }
        )
        result = int_brand_with_generic_available(df)
        assert len(result) == 0

    def test_empty_input(self) -> None:
        df = pl.DataFrame(
            {
                "product_rxcui": pl.Series([], dtype=pl.Utf8),
                "ndc_product_tty": pl.Series([], dtype=pl.Utf8),
                "product_startmarketingdate": pl.Series([], dtype=pl.Utf8),
                "ndc": pl.Series([], dtype=pl.Utf8),
            }
        )
        result = int_brand_with_generic_available(df)
        assert len(result) == 0


# ── int_claims_current_cost ──────────────────────────────────────────────────


class TestIntClaimsCurrentCost:
    def test_returns_dataframe(
        self,
        pharmacy_claim_df,
        ndc_terminology_df,
        rxnorm_brand_generic_df,
    ) -> None:
        ga = pl.DataFrame(
            {"brand_with_generic_available": pl.Series([], dtype=pl.Utf8)}
        )
        result = int_claims_current_cost(
            pharmacy_claim_df, ga, ndc_terminology_df, rxnorm_brand_generic_df
        )
        assert isinstance(result, pl.DataFrame)

    def test_aggregates_by_ndc(
        self,
        pharmacy_claim_df,
        ndc_terminology_df,
        rxnorm_brand_generic_df,
    ) -> None:
        ga = pl.DataFrame(
            {"brand_with_generic_available": pl.Series([], dtype=pl.Utf8)}
        )
        result = int_claims_current_cost(
            pharmacy_claim_df, ga, ndc_terminology_df, rxnorm_brand_generic_df
        )
        # NDC001 appears twice in claims, should aggregate
        ndc001 = result.filter(pl.col("ndc_code") == "NDC001")
        if len(ndc001) > 0:
            assert ndc001["paid_amount"][0] == pytest.approx(25.0)
            assert ndc001["claim_count"][0] == 2

    def test_computes_cost_per_claim(
        self,
        pharmacy_claim_df,
        ndc_terminology_df,
        rxnorm_brand_generic_df,
    ) -> None:
        ga = pl.DataFrame(
            {"brand_with_generic_available": pl.Series([], dtype=pl.Utf8)}
        )
        result = int_claims_current_cost(
            pharmacy_claim_df, ga, ndc_terminology_df, rxnorm_brand_generic_df
        )
        assert "cost_per_claim" in result.columns
        assert "cost_per_unit" in result.columns

    def test_expected_columns(
        self,
        pharmacy_claim_df,
        ndc_terminology_df,
        rxnorm_brand_generic_df,
    ) -> None:
        ga = pl.DataFrame(
            {"brand_with_generic_available": pl.Series([], dtype=pl.Utf8)}
        )
        result = int_claims_current_cost(
            pharmacy_claim_df, ga, ndc_terminology_df, rxnorm_brand_generic_df
        )
        for col in [
            "ndc_code",
            "paid_amount",
            "claim_count",
            "total_units",
            "cost_per_claim",
            "cost_per_unit",
        ]:
            assert col in result.columns


# ── int_generic_cost ─────────────────────────────────────────────────────────


class TestIntGenericCost:
    def test_filters_to_generic(self) -> None:
        claims = pl.DataFrame(
            {
                "ndc_code": ["NDC001", "NDC003"],
                "rxcui": ["RX001", "RX002"],
                "ndc_description": ["Drug A", "Drug B"],
                "data_source": ["test", "test"],
                "paid_amount": [10.0, 50.0],
                "claim_count": [2, 1],
                "total_units": [60.0, 30.0],
                "cost_per_claim": [5.0, 50.0],
                "cost_per_unit": [0.17, 1.67],
                "generic_available": [None, None],
                "brand_vs_generic": ["generic", "brand"],
            }
        )
        rxnorm = pl.DataFrame(
            {
                "product_rxcui": ["RX001"],
                "ndc": ["NDC001"],
            }
        )
        result = int_generic_cost(claims, rxnorm)
        assert isinstance(result, pl.DataFrame)
        assert "brand_rxcui" in result.columns

    def test_expected_columns(self) -> None:
        claims = pl.DataFrame(
            {
                "ndc_code": ["NDC001"],
                "rxcui": ["RX001"],
                "ndc_description": ["Drug A"],
                "data_source": ["test"],
                "paid_amount": [10.0],
                "claim_count": [2],
                "total_units": [60.0],
                "cost_per_claim": [5.0],
                "cost_per_unit": [0.17],
                "generic_available": [None],
                "brand_vs_generic": ["generic"],
            }
        )
        rxnorm = pl.DataFrame(
            {
                "product_rxcui": ["RX001"],
                "ndc": ["NDC001"],
            }
        )
        result = int_generic_cost(claims, rxnorm)
        for col in [
            "brand_rxcui",
            "data_source",
            "prescribed_atleast_one_generic_history",
            "generic_average_cost_per_unit",
        ]:
            assert col in result.columns


# ── int_generic_cost_by_ndc ──────────────────────────────────────────────────


class TestIntGenericCostByNdc:
    def test_filters_to_generic(self) -> None:
        claims = pl.DataFrame(
            {
                "ndc_code": ["NDC001", "NDC003"],
                "rxcui": ["RX001", "RX002"],
                "ndc_description": ["Drug A", "Drug B"],
                "data_source": ["test", "test"],
                "paid_amount": [10.0, 50.0],
                "claim_count": [2, 1],
                "total_units": [60.0, 30.0],
                "cost_per_claim": [5.0, 50.0],
                "cost_per_unit": [0.17, 1.67],
                "generic_available": [None, None],
                "brand_vs_generic": ["generic", "brand"],
            }
        )
        rxnorm = pl.DataFrame(
            {
                "product_rxcui": ["RX001"],
                "ndc": ["NDC001"],
            }
        )
        result = int_generic_cost_by_ndc(claims, rxnorm)
        assert isinstance(result, pl.DataFrame)
        assert len(result) >= 1

    def test_expected_columns(self) -> None:
        claims = pl.DataFrame(
            {
                "ndc_code": ["NDC001"],
                "rxcui": ["RX001"],
                "ndc_description": ["Drug A"],
                "data_source": ["test"],
                "paid_amount": [10.0],
                "claim_count": [2],
                "total_units": [60.0],
                "cost_per_claim": [5.0],
                "cost_per_unit": [0.17],
                "generic_available": [None],
                "brand_vs_generic": ["generic"],
            }
        )
        rxnorm = pl.DataFrame(
            {
                "product_rxcui": ["RX001"],
                "ndc": ["NDC001"],
            }
        )
        result = int_generic_cost_by_ndc(claims, rxnorm)
        for col in [
            "brand_rxcui",
            "generic_ndc_code",
            "paid_amount",
            "claim_count",
            "cost_per_unit",
            "prescribed_atleast_one_generic_history",
        ]:
            assert col in result.columns