stack/tests/aco/test_express_input_layer.py

"""Tests for aco.express.input_layer — schema definitions and passthrough functions."""

from __future__ import annotations

import polars as pl

from aco.express.input_layer import (
    condition,
    encounter,
    input_layer__appointment,
    input_layer__condition,
    input_layer__eligibility,
    input_layer__encounter,
    input_layer__immunization,
    input_layer__lab_result,
    input_layer__location,
    input_layer__medical_claim,
    input_layer__medication,
    input_layer__observation,
    input_layer__patient,
    input_layer__pharmacy_claim,
    input_layer__practitioner,
    input_layer__procedure,
    input_layer__provider_attribution,
    location,
    medication,
    patient,
    practitioner,
    procedure,
)

# ── schema definition functions ──────────────────────────────────────────────


class TestSchemaFunctions:
    """Schema functions project df to all-NULL columns with the right names."""

    def test_condition_schema(self) -> None:
        df = pl.DataFrame({"x": [1]})
        result = condition(df)
        assert isinstance(result, pl.DataFrame)
        assert len(result) == 1
        assert "condition_id" in result.columns
        assert "data_source" in result.columns
        assert result["condition_id"][0] is None

    def test_encounter_schema(self) -> None:
        df = pl.DataFrame({"x": [1]})
        result = encounter(df)
        assert "encounter_id" in result.columns
        assert "paid_amount" in result.columns
        assert len(result.columns) == 29

    def test_location_schema(self) -> None:
        df = pl.DataFrame({"x": [1]})
        result = location(df)
        assert "location_id" in result.columns
        assert "latitude" in result.columns
        assert len(result.columns) == 12

    def test_medication_schema(self) -> None:
        df = pl.DataFrame({"x": [1]})
        result = medication(df)
        assert "medication_id" in result.columns
        assert "ndc_code" in result.columns

    def test_patient_schema(self) -> None:
        df = pl.DataFrame({"x": [1]})
        result = patient(df)
        assert "person_id" in result.columns
        assert "birth_date" in result.columns

    def test_practitioner_schema(self) -> None:
        df = pl.DataFrame({"x": [1]})
        result = practitioner(df)
        assert "practitioner_id" in result.columns
        assert "npi" in result.columns
        assert len(result.columns) == 8

    def test_procedure_schema(self) -> None:
        df = pl.DataFrame({"x": [1]})
        result = procedure(df)
        assert "procedure_id" in result.columns
        assert "normalized_code" in result.columns

    def test_schema_preserves_row_count(self) -> None:
        df = pl.DataFrame({"x": [1]})
        assert len(condition(df)) == 1
        assert len(encounter(df)) == 1
        assert len(location(df)) == 1


# ── passthrough functions ────────────────────────────────────────────────────


class TestPassthroughFunctions:
    """Passthrough functions return input unchanged."""

    def test_condition_passthrough(self, condition_df) -> None:
        result = input_layer__condition(condition_df)
        assert result.shape == condition_df.shape
        assert result.columns == condition_df.columns

    def test_eligibility_passthrough(self, eligibility_df) -> None:
        result = input_layer__eligibility(eligibility_df)
        assert result.shape == eligibility_df.shape

    def test_encounter_passthrough(self, encounter_df) -> None:
        result = input_layer__encounter(encounter_df)
        assert result.shape == encounter_df.shape

    def test_location_passthrough(self, location_df) -> None:
        result = input_layer__location(location_df)
        assert result.shape == location_df.shape

    def test_medical_claim_passthrough(self, medical_claim_df) -> None:
        result = input_layer__medical_claim(medical_claim_df)
        assert result.shape == medical_claim_df.shape

    def test_medication_passthrough(self, medication_df) -> None:
        result = input_layer__medication(medication_df)
        assert result.shape == medication_df.shape

    def test_observation_passthrough(self, observation_df) -> None:
        result = input_layer__observation(observation_df)
        assert result.shape == observation_df.shape

    def test_patient_passthrough(self, patient_df) -> None:
        result = input_layer__patient(patient_df)
        assert result.shape == patient_df.shape

    def test_pharmacy_claim_passthrough(self, pharmacy_claim_df) -> None:
        result = input_layer__pharmacy_claim(pharmacy_claim_df)
        assert result.shape == pharmacy_claim_df.shape

    def test_practitioner_passthrough(self, practitioner_df) -> None:
        result = input_layer__practitioner(practitioner_df)
        assert result.shape == practitioner_df.shape

    def test_procedure_passthrough(self, procedure_df) -> None:
        result = input_layer__procedure(procedure_df)
        assert result.shape == procedure_df.shape

    def test_provider_attribution_passthrough(self, provider_attribution_df) -> None:
        result = input_layer__provider_attribution(provider_attribution_df)
        assert result.shape == provider_attribution_df.shape

    def test_passthrough_preserves_values(self, patient_df) -> None:
        result = input_layer__patient(patient_df)
        assert result["first_name"][0] == "Jane"
        assert result["person_id"][0] == "P001"

    def test_appointment_passthrough(self, appointment_df) -> None:
        result = input_layer__appointment(appointment_df)
        assert result.shape == appointment_df.shape

    def test_immunization_passthrough(self, immunization_df) -> None:
        result = input_layer__immunization(immunization_df)
        assert result.shape == immunization_df.shape

    def test_lab_result_passthrough(self, lab_result_df) -> None:
        result = input_layer__lab_result(lab_result_df)
        assert result.shape == lab_result_df.shape

    def test_empty_passthrough(self) -> None:
        df = pl.DataFrame({"x": pl.Series([], dtype=pl.Utf8)})
        result = input_layer__procedure(df)
        assert len(result) == 0