stack/tests/opps/test_opps_calcs.py

"""Tests for opps.calcs.payment — OPPS payment calculations."""

from __future__ import annotations

import polars as pl

from opps.calcs.payment import payment, skin_sub_impact

# ── Fixtures ──────────────────────────────────────────────────────────────────


def _apc_df():
    return pl.DataFrame(
        {
            "hcpcs": ["G0499", "G0500"],
            "apc": ["5115", "5116"],
            "relative_weight": [1.2, 0.8],
            "payment_rate": [100.0, 80.0],
            "status_indicator": ["J1", "J1"],
            "cbsa": ["10180", "99999"],
        }
    )


def _wage_index_df():
    return pl.DataFrame(
        {
            "cbsa": ["10180"],
            "wage_index": [1.05],
        }
    )


def _claims_df():
    return pl.DataFrame(
        {
            "hcpcs_code": ["Q4151", "Q4152"],
            "units": [2, 3],
            "paid_amount": [250.0, 400.0],
        }
    )


def _asp_df():
    return pl.DataFrame(
        {
            "hcpcs_code": ["Q4151", "Q4152"],
            "asp_per_unit": [100.0, 120.0],
            "payment_limit": [106.0, 127.2],  # ASP + 6%
        }
    )


# ── payment() ─────────────────────────────────────────────────────────────────


class TestPayment:
    def test_returns_dataframe(self) -> None:
        result = payment(_apc_df(), _wage_index_df())
        assert isinstance(result, pl.DataFrame)

    def test_adjusted_payment_column_present(self) -> None:
        result = payment(_apc_df(), _wage_index_df())
        assert "adjusted_payment" in result.columns

    def test_row_count_matches_input(self) -> None:
        result = payment(_apc_df(), _wage_index_df())
        assert len(result) == 2

    def test_payment_formula_with_wage_index(self) -> None:
        """For cbsa 10180: payment = 100 * (0.60 * 1.05 + 0.40) = 103.0."""
        result = payment(_apc_df(), _wage_index_df())
        row = result.filter(pl.col("hcpcs") == "G0499")
        expected = round(100.0 * (0.60 * 1.05 + 0.40), 2)
        assert abs(row["adjusted_payment"][0] - expected) < 0.01

    def test_missing_cbsa_uses_default_wage_index(self) -> None:
        """cbsa 99999 not in wage_index → fill_null(1.0) → no adjustment."""
        result = payment(_apc_df(), _wage_index_df())
        row = result.filter(pl.col("hcpcs") == "G0500")
        expected = round(80.0 * (0.60 * 1.0 + 0.40), 2)
        assert abs(row["adjusted_payment"][0] - expected) < 0.01

    def test_custom_labor_share(self) -> None:
        result = payment(_apc_df(), _wage_index_df(), labor_share=0.50)
        row = result.filter(pl.col("hcpcs") == "G0499")
        expected = round(100.0 * (0.50 * 1.05 + 0.50), 2)
        assert abs(row["adjusted_payment"][0] - expected) < 0.01


# ── skin_sub_impact() ─────────────────────────────────────────────────────────


class TestSkinSubImpact:
    def test_returns_dataframe(self) -> None:
        result = skin_sub_impact(_claims_df(), _asp_df())
        assert isinstance(result, pl.DataFrame)

    def test_expected_columns(self) -> None:
        result = skin_sub_impact(_claims_df(), _asp_df())
        for col in ("old_payment", "new_payment", "payment_delta", "pct_change"):
            assert col in result.columns

    def test_old_payment_formula(self) -> None:
        """old_payment = payment_limit * units."""
        result = skin_sub_impact(_claims_df(), _asp_df())
        row = result.filter(pl.col("hcpcs_code") == "Q4151")
        assert row["old_payment"][0] == round(106.0 * 2, 2)

    def test_new_payment_formula(self) -> None:
        """new_payment = flat_rate * units (default $127.28)."""
        result = skin_sub_impact(_claims_df(), _asp_df())
        row = result.filter(pl.col("hcpcs_code") == "Q4151")
        assert row["new_payment"][0] == round(127.28 * 2, 2)

    def test_payment_delta(self) -> None:
        result = skin_sub_impact(_claims_df(), _asp_df())
        row = result.filter(pl.col("hcpcs_code") == "Q4151")
        old = round(106.0 * 2, 2)
        new = round(127.28 * 2, 2)
        assert abs(row["payment_delta"][0] - round(new - old, 2)) < 0.01

    def test_pct_change_computed(self) -> None:
        result = skin_sub_impact(_claims_df(), _asp_df())
        row = result.filter(pl.col("hcpcs_code") == "Q4151")
        assert row["pct_change"][0] is not None

    def test_custom_flat_rate(self) -> None:
        result = skin_sub_impact(_claims_df(), _asp_df(), flat_rate=200.0)
        row = result.filter(pl.col("hcpcs_code") == "Q4151")
        assert row["new_payment"][0] == round(200.0 * 2, 2)

    def test_row_count_matches_input(self) -> None:
        result = skin_sub_impact(_claims_df(), _asp_df())
        assert len(result) == 2


# ── opps.pipe — STEPS list ────────────────────────────────────────────────────


class TestOppsPipe:
    def test_steps_is_list(self) -> None:
        from opps.pipe import STEPS

        assert isinstance(STEPS, list)

    def test_steps_populated(self) -> None:
        from opps.pipe import STEPS

        assert len(STEPS) == 2

    def test_step_names(self) -> None:
        from opps.pipe import STEPS

        names = [s[0] for s in STEPS]
        assert "opps_adjusted_payment" in names
        assert "opps_skin_sub_impact" in names

    def test_steps_have_output_tables(self) -> None:
        from opps.pipe import STEPS
        from opps.table import OppsAdjustedPayment, OppsSkinSubImpact

        outputs = {s[0]: s[2] for s in STEPS}
        assert outputs["opps_adjusted_payment"] is OppsAdjustedPayment
        assert outputs["opps_skin_sub_impact"] is OppsSkinSubImpact


# ── Express functions ─────────────────────────────────────────────────────────


class TestExpressAdjustedPayment:
    def test_import(self) -> None:
        from opps.express.payment import adjusted_payment

        assert callable(adjusted_payment)

    def test_runs_with_dataframes(self) -> None:
        from opps.express.payment import adjusted_payment

        result = adjusted_payment(_apc_df(), _wage_index_df())
        assert "adjusted_payment" in result.columns
        assert len(result) == 2

    def test_formula_matches_calcs(self) -> None:
        from opps.express.payment import adjusted_payment

        direct = payment(_apc_df(), _wage_index_df())
        via_express = adjusted_payment(_apc_df(), _wage_index_df())
        assert (
            direct["adjusted_payment"].to_list()
            == via_express["adjusted_payment"].to_list()
        )


class TestExpressSkinSubImpact:
    def test_import(self) -> None:
        from opps.express.payment import skin_sub_impact

        assert callable(skin_sub_impact)

    def test_runs_with_pass_through_and_asp(self) -> None:
        from opps.express.payment import skin_sub_impact as express_fn

        pass_through = pl.DataFrame(
            {
                "hcpcs": ["Q4151", "Q4152"],
                "short_description": ["Skin sub A", "Skin sub B"],
                "pass_through_type": ["biological", "biological"],
                "status_indicator": ["G", "G"],
                "payment_rate": [250.0, 400.0],
                "asp_per_unit": [100.0, 120.0],
                "year": [2025, 2025],
            }
        )
        asp = pl.DataFrame(
            {
                "hcpcs_code": ["Q4151", "Q4152"],
                "asp_per_unit": [100.0, 120.0],
            }
        )
        result = express_fn(pass_through, asp)
        assert "old_payment" in result.columns
        assert "new_payment" in result.columns
        assert len(result) == 2


# ── Pipeline registry ─────────────────────────────────────────────────────────


class TestOppsPipeline:
    def test_registered_in_registry(self) -> None:
        from aco.pipe import registry

        assert "opps" in registry

    def test_pipeline_has_two_exprs(self) -> None:
        from aco.pipe import registry

        assert len(registry["opps"]) == 2

    def test_pipeline_expr_names(self) -> None:
        from aco.pipe import registry

        names = registry["opps"].names()
        assert names == ["opps.adjusted_payment", "opps.skin_sub_impact"]

    def test_pipeline_inputs_resolved(self) -> None:
        from aco.pipe import registry

        p = registry["opps"]
        inputs = p.exprs[0].inputs
        assert "opps.apc_weight" in inputs
        assert "opps.wage_index" in inputs