app/test_deterministic_rank.py

"""Unit tests for clinic-pilot deterministic code selection."""
from __future__ import annotations

from .deterministic_rank import (
    build_deterministic_response,
    select_premium_docs,
    select_primary_docs,
)


def _doc(code: str, *, section: str = "PAEDIATRICS (26)", **extra) -> dict:
    return {
        "billing_code": code,
        "description_text": extra.pop("description_text", f"{code} service"),
        "base_fee_cad": 100.0,
        "parent_section": section,
        "in_current_schedule": True,
        "rules_and_constraints": "",
        **extra,
    }


def test_primary_order_stable_with_prior_boost():
    context = [
        _doc("A262"),
        _doc("A261"),
        _doc("A263"),
    ]
    # Without priors, retrieval order wins.
    a = [d["billing_code"] for d in select_primary_docs(context, limit=2)]
    assert a == ["A262", "A261"]
    # Strong prior on A261 flips it to rank 1 deterministically.
    b = [
        d["billing_code"]
        for d in select_primary_docs(
            context, learned_priors={"A261": 0.9}, limit=2
        )
    ]
    assert b == ["A261", "A262"]


def test_premium_skips_chronic_without_evidence():
    primaries = [
        _doc(
            "A261",
            relationship_premiums=["E078", "G372", "E080", "E402"],
        )
    ]
    pool = [
        _doc(
            "E078",
            section="CHRONIC DISEASE ASSESSMENT PREMIUM",
            description_text="chronic disease assessment premium",
        ),
        _doc(
            "G372",
            section="IMMUNIZATION",
            description_text="immunization injection",
        ),
        _doc(
            "E080",
            section="FIRST VISIT BY PRIMARY CARE PHYSICIAN AFTER HOSPITAL DISCHARGE",
            description_text="first visit after hospital discharge",
        ),
        _doc(
            "E402",
            section="AFTER HOURS SPECIAL VISIT",
            description_text="evenings Monday to Friday and weekends/holidays",
        ),
    ]
    # Daytime acute visit — no chronic/immuniz/discharge/time premium.
    picked = select_premium_docs(
        pool,
        primary_docs=primaries,
        clinical_summary="Healthy 8yo with acute otitis media, brief office visit.",
        time_align=None,
        age_align=None,
        limit=5,
    )
    codes = [d["billing_code"] for d in picked]
    assert "E078" not in codes
    assert "G372" not in codes
    assert "E080" not in codes


def test_premium_includes_chronic_when_note_supports():
    primaries = [_doc("A261", relationship_premiums=["E078"])]
    pool = [
        _doc(
            "E078",
            section="CHRONIC DISEASE ASSESSMENT PREMIUM",
            description_text="chronic disease assessment premium",
        )
    ]
    picked = select_premium_docs(
        pool,
        primary_docs=primaries,
        clinical_summary="Follow-up for chronic asthma in a school-aged child.",
        time_align=None,
        age_align=None,
        limit=3,
    )
    assert [d["billing_code"] for d in picked] == ["E078"]


def test_paeds_boosts_a268_for_enhanced_18_month():
    context = [
        _doc("A261", description_text="Level 1 paediatric assessment"),
        _doc("A268", description_text="Enhanced 18 month well baby visit"),
        _doc("A262", description_text="Level 2 paediatric assessment"),
    ]
    picked = select_primary_docs(
        context,
        clinical_summary="Enhanced 18-month well-baby visit for toddler.",
        encounter_type="Ambulatory - Enhanced 18-Month Well-Baby Visit",
        age_months=18,
        limit=3,
    )
    assert [d["billing_code"] for d in picked][0] == "A268"


def test_paeds_drops_a268_without_18_month_evidence():
    context = [
        _doc("A261", description_text="Level 1 paediatric assessment"),
        _doc("A268", description_text="Enhanced 18 month well baby visit"),
        _doc("A262", description_text="Level 2 paediatric assessment"),
    ]
    picked = select_primary_docs(
        context,
        clinical_summary="Otherwise healthy 8-year-old with acute otitis media.",
        encounter_type="Ambulatory - Level 1 Paediatric Assessment",
        age_months=96,
        prefixes=["A"],
        limit=3,
    )
    codes = [d["billing_code"] for d in picked]
    assert codes[0] == "A261"
    assert "A268" not in codes


def test_prefix_filter_drops_wrong_setting():
    context = [
        _doc("C260", description_text="Hospital special paediatric consultation"),
        _doc("A261", description_text="Level 1 paediatric assessment"),
    ]
    picked = select_primary_docs(
        context,
        clinical_summary="Brief office visit for otitis.",
        encounter_type="Ambulatory - Level 1 Paediatric Assessment",
        prefixes=["A"],
        limit=3,
    )
    codes = [d["billing_code"] for d in picked]
    assert codes == ["A261"]


def test_template_justification_explains_code_without_llm_jargon():
    primary = [
        _doc(
            "A261",
            description_text="Level 1 paediatric assessment",
            parent_section="PAEDIATRICS (26)",
            reference="A20",
            rules_and_constraints="Limited assessment of a single system or complaint.",
            relationship_family="office_assessment",
            relationship_ladder=(
                "Office ladder: Level 1 A261 → Level 2 A262 → medical specific A263"
            ),
        )
    ]
    result = build_deterministic_response(
        primary,
        [],
        clinical_summary="Otherwise healthy 8-year-old with acute otitis media.",
        encounter_type="Ambulatory - Level 1 Paediatric Assessment",
        specialty_label="Paediatrics",
    )
    text = result.top_matching_codes[0].justification
    assert "A261" in text
    assert "Level 1" in text
    assert "otitis" in text.lower() or "acute" in text.lower()
    assert "retrieval rank" not in text.lower()
    assert "PAEDIATRICS" in text
    assert "Constraint:" in text or "Limited assessment" in text