stack/tests/aco/test_express_claims_preprocessing.py

"""Tests for aco.express.claims_preprocessing — encounter grouping functions.

Tests verify the decorated wrapper functions and the encounter type chains
(anchor_events → generate_encounter_id → match_claims_to_anchor).
The _generate_encounter_id and _match_claims_to_anchor helpers are tested
indirectly through their decorated callers since they consume narwhals frames.
"""

from __future__ import annotations

from datetime import date

import polars as pl
import pytest

from aco.express.claims_preprocessing import (
    ambulance__anchor_events,
    ambulance__generate_encounter_id,
    ambulance__match_claims_to_anchor,
    asc__anchor_events,
    asc__generate_encounter_id,
    asc__start_end_dates,
    dialysis__anchor_events,
    dialysis__generate_encounter_id,
    dialysis__match_claims_to_anchor,
)

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


@pytest.fixture
def stg_medical_claim() -> pl.DataFrame:
    """Staged medical claims for encounter grouping."""
    return pl.DataFrame(
        {
            "claim_id": ["C1", "C2", "C3", "C4"],
            "claim_line_number": [1, 1, 1, 1],
            "patient_data_source_id": ["PDS1", "PDS1", "PDS2", "PDS1"],
            "start_date": [
                date(2024, 1, 15),
                date(2024, 1, 15),
                date(2024, 2, 1),
                date(2024, 3, 1),
            ],
            "end_date": [
                date(2024, 1, 20),
                date(2024, 1, 20),
                date(2024, 2, 5),
                date(2024, 3, 5),
            ],
        }
    )


# ── ambulance chain ──────────────────────────────────────────────────────────


class TestAmbulanceChain:
    def test_anchor_events_selects_claim_id(self, stg_medical_claim) -> None:
        result = ambulance__anchor_events(stg_medical_claim)
        assert result.columns == ["claim_id"]
        assert len(result) == len(stg_medical_claim)

    def test_generate_encounter_id(self, stg_medical_claim) -> None:
        anchor = ambulance__anchor_events(stg_medical_claim)
        result = ambulance__generate_encounter_id(anchor, stg_medical_claim)
        assert isinstance(result, pl.DataFrame)
        assert "old_encounter_id" in result.columns

    def test_same_patient_date_same_encounter(self, stg_medical_claim) -> None:
        anchor = ambulance__anchor_events(stg_medical_claim)
        result = ambulance__generate_encounter_id(anchor, stg_medical_claim)
        # C1 and C2 have same patient+date → same encounter_id
        ids = result.sort("claim_id")["old_encounter_id"].to_list()
        assert ids[0] == ids[1]

    def test_different_patients_different_encounters(self, stg_medical_claim) -> None:
        anchor = ambulance__anchor_events(stg_medical_claim)
        result = ambulance__generate_encounter_id(anchor, stg_medical_claim)
        enc_ids = result["old_encounter_id"].unique()
        assert len(enc_ids) >= 2

    def test_expected_columns_generate(self, stg_medical_claim) -> None:
        anchor = ambulance__anchor_events(stg_medical_claim)
        result = ambulance__generate_encounter_id(anchor, stg_medical_claim)
        for col in [
            "patient_data_source_id",
            "start_date",
            "claim_id",
            "old_encounter_id",
        ]:
            assert col in result.columns

    def test_match_claims(self, stg_medical_claim) -> None:
        anchor = ambulance__anchor_events(stg_medical_claim)
        gen = ambulance__generate_encounter_id(anchor, stg_medical_claim)
        result = ambulance__match_claims_to_anchor(gen, stg_medical_claim)
        assert "old_encounter_id" in result.columns
        assert len(result) > 0
        for col in [
            "patient_data_source_id",
            "start_date",
            "claim_id",
            "claim_line_number",
            "old_encounter_id",
        ]:
            assert col in result.columns

    def test_empty_claims(self) -> None:
        empty = pl.DataFrame(
            {
                "claim_id": pl.Series([], dtype=pl.Utf8),
                "claim_line_number": pl.Series([], dtype=pl.Int64),
                "patient_data_source_id": pl.Series([], dtype=pl.Utf8),
                "start_date": pl.Series([], dtype=pl.Date),
                "end_date": pl.Series([], dtype=pl.Date),
            }
        )
        anchor = ambulance__anchor_events(empty)
        assert len(anchor) == 0


# ── ASC chain ────────────────────────────────────────────────────────────────


class TestAscChain:
    def test_anchor_events(self, stg_medical_claim) -> None:
        result = asc__anchor_events(stg_medical_claim)
        assert result.columns == ["claim_id"]

    def test_generate_encounter_id(self, stg_medical_claim) -> None:
        anchor = asc__anchor_events(stg_medical_claim)
        result = asc__generate_encounter_id(anchor, stg_medical_claim)
        assert "old_encounter_id" in result.columns
        assert "end_date" in result.columns

    def test_start_end_dates(self, stg_medical_claim) -> None:
        anchor = asc__anchor_events(stg_medical_claim)
        gen = asc__generate_encounter_id(anchor, stg_medical_claim)
        result = asc__start_end_dates(gen)
        assert "encounter_start_date" in result.columns
        assert "encounter_end_date" in result.columns
        assert "patient_data_source_id" in result.columns

    def test_start_end_dates_aggregates(self, stg_medical_claim) -> None:
        anchor = asc__anchor_events(stg_medical_claim)
        gen = asc__generate_encounter_id(anchor, stg_medical_claim)
        result = asc__start_end_dates(gen)
        # Each unique patient+encounter group should have one row
        assert len(result) >= 1


# ── Dialysis chain ───────────────────────────────────────────────────────────


class TestDialysisChain:
    def test_anchor_events(self, stg_medical_claim) -> None:
        result = dialysis__anchor_events(stg_medical_claim)
        assert result.columns == ["claim_id"]

    def test_generate_encounter_id(self, stg_medical_claim) -> None:
        anchor = dialysis__anchor_events(stg_medical_claim)
        result = dialysis__generate_encounter_id(anchor, stg_medical_claim)
        assert "old_encounter_id" in result.columns

    def test_match_claims(self, stg_medical_claim) -> None:
        anchor = dialysis__anchor_events(stg_medical_claim)
        gen = dialysis__generate_encounter_id(anchor, stg_medical_claim)
        result = dialysis__match_claims_to_anchor(gen, stg_medical_claim)
        assert "old_encounter_id" in result.columns
        assert len(result) > 0