test_core.py

"""test_core.py — TDD tests for core.py (Drive & Save).

Run with: pytest test_core.py -v
"""
import h3
import pandas as pd
import pytest

import core
import data_source

# ---------------------------------------------------------------------------
# Part A — config loading
# ---------------------------------------------------------------------------


def test_load_config_reads_yaml():
    cfg = core.load_config()
    assert cfg["h3"]["resolution"] == 8
    assert cfg["distance"]["avg_speed_kmh"] > 0
    assert cfg["ranking"]["min_hospitals"] >= 1
    assert cfg["funnel"]["cta_base_url"].startswith("https://")


# ---------------------------------------------------------------------------
# Part B — H3 distance with haversine fallback
# ---------------------------------------------------------------------------

CFG = core.load_config()
BOULDER = h3.latlng_to_cell(40.0150, -105.2705, 8)
DENVER = h3.latlng_to_cell(39.7392, -104.9903, 8)


def test_grid_distance_km_same_cell_is_zero():
    assert core.grid_distance_km(DENVER, DENVER, CFG) == 0.0


def test_grid_distance_km_boulder_denver_is_realistic():
    km = core.grid_distance_km(BOULDER, DENVER, CFG)
    assert 30 <= km <= 50


def test_grid_distance_km_falls_back_on_h3_error(monkeypatch):
    def boom(*args, **kwargs):
        raise ValueError("cells too far apart")
    monkeypatch.setattr(h3, "grid_distance", boom)
    km = core.grid_distance_km(BOULDER, DENVER, CFG)
    assert 30 <= km <= 50


# ---------------------------------------------------------------------------
# Part C — drive-time estimate
# ---------------------------------------------------------------------------


def test_drive_minutes_zero_distance():
    assert core.drive_minutes(0.0, CFG) == 0


def test_drive_minutes_rounds_to_nearest_five():
    m = core.drive_minutes(39.0, CFG)
    assert m % 5 == 0
    assert 30 <= m <= 50


def test_drive_minutes_monotonic():
    assert core.drive_minutes(80.0, CFG) > core.drive_minutes(20.0, CFG)


# ---------------------------------------------------------------------------
# Part D — ranking, suppression, savings
# ---------------------------------------------------------------------------

SPRINGS = h3.latlng_to_cell(38.8339, -104.8214, 8)


def _fixtures():
    medians = pd.DataFrame([
        {"state": "CO", "procedure_code": "45378", "procedure_name": "Diagnostic colonoscopy",
         "metro": "Denver", "median_price": 2267, "n_hospitals": 12},
        {"state": "CO", "procedure_code": "45378", "procedure_name": "Diagnostic colonoscopy",
         "metro": "Boulder", "median_price": 335, "n_hospitals": 5},
        {"state": "CO", "procedure_code": "45378", "procedure_name": "Diagnostic colonoscopy",
         "metro": "Colorado Springs", "median_price": 1122, "n_hospitals": 7},
        {"state": "CO", "procedure_code": "45378", "procedure_name": "Diagnostic colonoscopy",
         "metro": "Tiny Town", "median_price": 50, "n_hospitals": 1},
    ])
    metros = pd.DataFrame([
        {"state": "CO", "metro": "Denver", "h3_cell": DENVER},
        {"state": "CO", "metro": "Boulder", "h3_cell": BOULDER},
        {"state": "CO", "metro": "Colorado Springs", "h3_cell": SPRINGS},
        {"state": "CO", "metro": "Tiny Town", "h3_cell": SPRINGS},
    ])
    return medians, metros


def test_rank_metros_picks_cheapest_non_suppressed():
    medians, metros = _fixtures()
    r = core.rank_metros(medians, metros, "45378", "Denver", CFG)
    assert r.home_metro == "Denver"
    assert r.home_median == 2267
    assert r.cheapest_metro == "Boulder"
    assert r.cheapest_median == 335
    assert r.savings == 2267 - 335
    assert round(r.savings_pct) == 85
    assert r.drive_min_to_cheapest >= 30
    assert r.home_is_cheapest is False


def test_rank_metros_suppresses_low_n():
    medians, metros = _fixtures()
    r = core.rank_metros(medians, metros, "45378", "Denver", CFG)
    visible = [row.metro for row in r.rows if not row.suppressed]
    assert "Tiny Town" not in visible
    assert {"Denver", "Boulder", "Colorado Springs"}.issubset(set(visible))


def test_rank_metros_home_is_cheapest_branch():
    medians, metros = _fixtures()
    r = core.rank_metros(medians, metros, "45378", "Boulder", CFG)
    assert r.home_is_cheapest is True
    assert r.cheapest_metro == "Colorado Springs"
    assert r.savings == 0


# ---------------------------------------------------------------------------
# Part E — CTA URL builder
# ---------------------------------------------------------------------------


def test_build_cta_url_homepage_fallback_when_no_pattern():
    url = core.build_cta_url("45378", "Boulder", CFG)
    assert url.startswith("https://clearprice-health.com")
    assert "utm_source=mile-hi-labs" in url


def test_rank_metros_home_suppressed_no_negative_savings():
    medians, metros = _fixtures()
    # "Tiny Town" has n_hospitals=1 (suppressed) and the lowest price ($50)
    r = core.rank_metros(medians, metros, "45378", "Tiny Town", CFG)
    assert r.savings >= 0            # never negative
    assert r.savings == 0            # home price unreliable -> no claimed savings
    assert r.home_is_cheapest is False
    assert r.cheapest_metro == "Boulder"   # cheapest non-suppressed shown informationally


def test_build_cta_url_uses_deep_link_pattern():
    cfg = core.load_config()
    cfg["funnel"]["deep_link_pattern"] = "/procedure/{code}/?metro={metro_slug}"
    url = core.build_cta_url("45378", "Colorado Springs", cfg)
    assert "/procedure/45378/" in url
    assert "metro=colorado-springs" in url
    assert "utm_source=mile-hi-labs" in url


# ---------------------------------------------------------------------------
# Part G — app.py: import contract + format_headline
# ---------------------------------------------------------------------------

import app as app_module


def test_app_imports_and_exposes_build_demo():
    assert hasattr(app_module, "build_demo")


def test_format_headline_savings():
    medians, metros = _fixtures()
    r = core.rank_metros(medians, metros, "45378", "Denver", CFG)
    text = app_module.format_headline(r, CFG)
    assert "Denver" in text and "Boulder" in text
    assert "$1,932" in text
    assert "~" in text  # approx drive time label


def test_format_headline_no_comparable_metro():
    import pandas as pd
    medians2 = pd.DataFrame([
        {"state": "CO", "procedure_code": "99999", "procedure_name": "Solo proc",
         "metro": "Denver", "median_price": 500, "n_hospitals": 12},
        {"state": "CO", "procedure_code": "99999", "procedure_name": "Solo proc",
         "metro": "Tiny Town", "median_price": 100, "n_hospitals": 1},
    ])
    _, metros = _fixtures()
    r = core.rank_metros(medians2, metros, "99999", "Denver", CFG)
    text = app_module.format_headline(r, CFG)   # must not raise
    assert "Denver" in text
    assert "$500" in text
    assert "None" not in text


# ---------------------------------------------------------------------------
# Part F — data_source: load curated slice from private HF Dataset
# ---------------------------------------------------------------------------


def test_load_slices_uses_token_and_reads_csvs(monkeypatch, tmp_path):
    medians_csv = tmp_path / "medians.csv"
    metros_csv = tmp_path / "metros.csv"
    medians_csv.write_text(
        "state,procedure_code,procedure_name,metro,median_price,n_hospitals\n"
        "CO,45378,Diagnostic colonoscopy,Denver,2267,12\n"
    )
    metros_csv.write_text("state,metro,h3_cell\nCO,Denver,8a2a1072b59ffff\n")

    calls = {}
    def fake_download(repo_id, filename, repo_type, token):
        calls["repo_id"] = repo_id
        calls["repo_type"] = repo_type
        calls["token"] = token
        return str(medians_csv if filename.endswith("medians.csv") else metros_csv)

    monkeypatch.setattr(data_source, "hf_hub_download", fake_download)
    cfg = core.load_config()
    medians, metros = data_source.load_slices(cfg, token="hf_test")

    assert calls["repo_type"] == "dataset"
    assert calls["token"] == "hf_test"
    assert calls["repo_id"] == cfg["data"]["dataset_repo"]
    assert list(medians["metro"]) == ["Denver"]
    assert list(metros["metro"]) == ["Denver"]


# ---------------------------------------------------------------------------
# Part H — app builds under gradio 5
# ---------------------------------------------------------------------------
import app as app_module

def test_build_demo_returns_blocks():
    import gradio as gr
    demo = app_module.build_demo()
    assert isinstance(demo, gr.Blocks)


# ---------------------------------------------------------------------------
# Part I — llm config + module
# ---------------------------------------------------------------------------

def test_config_has_llm_block():
    cfg = core.load_config()
    assert cfg["llm"]["zerogpu_model_id"]
    assert cfg["llm"]["recommendation_max_sentences"] >= 1
    assert "not medical advice" in cfg["llm"]["medical_disclaimer"]


import llm

def _fake_gen(returns):
    return lambda system, user: returns

CATALOG = {
    "procedures": {"Diagnostic colonoscopy": "45378", "MRI brain": "70551"},
    "metros": ["Denver Metro", "Boulder", "Colorado Springs"],
}

def test_parse_ok_maps_name_to_code():
    gen = _fake_gen('{"procedure": "Diagnostic colonoscopy", "metro": "Denver Metro"}')
    r = llm.parse_request("cheapest colonoscopy near denver", CATALOG, gen)
    assert r.status == "ok"
    assert r.procedure_code == "45378"
    assert r.metro == "Denver Metro"

def test_parse_missing_metro():
    gen = _fake_gen('{"procedure": "MRI brain", "metro": null}')
    r = llm.parse_request("mri brain", CATALOG, gen)
    assert r.status == "missing_metro"
    assert r.procedure_code == "70551" and r.metro is None

def test_parse_missing_procedure():
    gen = _fake_gen('{"procedure": null, "metro": "Boulder"}')
    r = llm.parse_request("prices in boulder", CATALOG, gen)
    assert r.status == "missing_procedure"

def test_parse_hallucinated_values_rejected():
    gen = _fake_gen('{"procedure": "Brain transplant", "metro": "Atlantis"}')
    r = llm.parse_request("brain transplant in atlantis", CATALOG, gen)
    assert r.status == "no_match"
    assert r.procedure_code is None and r.metro is None

def test_parse_bad_json():
    gen = _fake_gen("I think you want a colonoscopy in Denver, friend!")
    r = llm.parse_request("...", CATALOG, gen)
    assert r.status == "bad_json"


def _denver_ranking():
    medians, metros = _fixtures()
    return core.rank_metros(medians, metros, "45378", "Denver", CFG)

def test_recommendation_passes_through_grounded_prose():
    r = _denver_ranking()
    gen = _fake_gen("Drive to Boulder and pay $335 instead of $2,267 — save $1,932 (85%).")
    out = llm.write_recommendation(r, CFG, fallback_text="FALLBACK", generate=gen)
    assert "Boulder" in out
    assert "FALLBACK" not in out
    assert CFG["llm"]["medical_disclaimer"] in out

def test_recommendation_with_ungrounded_number_falls_back():
    r = _denver_ranking()
    gen = _fake_gen("You could save up to $9,999 by driving to Boulder!")
    out = llm.write_recommendation(r, CFG, fallback_text="FALLBACK", generate=gen)
    assert out == "FALLBACK"


def test_recommendation_rejects_bare_ungrounded_number():
    r = _denver_ranking()
    gen = _fake_gen("You could save 9999 dollars driving to Boulder.")
    out = llm.write_recommendation(r, CFG, fallback_text="FALLBACK", generate=gen)
    assert out == "FALLBACK"


def test_import_safe_without_torch():
    import importlib
    importlib.reload(llm)
    assert llm.zerogpu_available() in (True, False)  # no ImportError raised

def test_detect_provider_prefers_explicit():
    assert llm.detect_provider({"MODEL_PROVIDER": "huggingface"}) == "huggingface"

def test_detect_provider_space_with_deps_is_zerogpu(monkeypatch):
    monkeypatch.setattr(llm, "zerogpu_available", lambda: True)
    assert llm.detect_provider({"SPACE_ID": "x/y"}) == "zerogpu"

def test_detect_provider_space_without_deps_is_huggingface(monkeypatch):
    monkeypatch.setattr(llm, "zerogpu_available", lambda: False)
    assert llm.detect_provider({"SPACE_ID": "x/y"}) == "huggingface"


# ---------------------------------------------------------------------------
# Part J — app.py: result builders (Task 6)
# ---------------------------------------------------------------------------

def test_build_catalog_shape():
    medians, metros = _fixtures()
    cat = app_module.build_catalog(medians, metros)
    names = {p["name"]: p["code"] for p in cat["procedures"]}
    assert names.get("Diagnostic colonoscopy") == "45378"
    assert any(m["label"] == "Denver" for m in cat["metros"])

def test_build_lookup_result():
    medians, metros = _fixtures()
    res = app_module.build_lookup(medians, metros, "45378", "Denver", CFG)
    assert res["status"] == "ok"
    assert res["recommendation"] is None
    assert res["headline"] and "Boulder" in res["headline"]
    metros_in_rows = {row["metro"] for row in res["rows"]}
    assert "Tiny Town" not in metros_in_rows
    boulder = next(r for r in res["rows"] if r["metro"] == "Boulder")
    assert boulder["is_cheapest"] is True
    assert res["cta_url"].startswith("https://")

def test_build_ask_ok_includes_recommendation():
    medians, metros = _fixtures()
    calls = {"n": 0}
    def gen2(system, user):
        calls["n"] += 1
        if calls["n"] == 1:
            return '{"procedure": "Diagnostic colonoscopy", "metro": "Denver"}'
        return "Drive to Boulder for $335 instead of $2,267 — save $1,932 (85%)."
    res = app_module.build_ask(medians, metros, "cheapest colonoscopy in denver", CFG, gen2)
    assert res["status"] == "ok"
    assert res["interpretation"]["metro"] == "Denver"
    assert "Boulder" in res["recommendation"]

def test_build_ask_missing_metro_status():
    medians, metros = _fixtures()
    gen = _fake_gen('{"procedure": "Diagnostic colonoscopy", "metro": null}')
    res = app_module.build_ask(medians, metros, "colonoscopy", CFG, gen)
    assert res["status"] == "missing_metro"
    assert res["rows"] == [] and res["recommendation"] is None


# ---------------------------------------------------------------------------
# Part J — API endpoint reachability (the API seam)
# ---------------------------------------------------------------------------

def test_named_endpoints_reachable_via_client():
    import gradio as gr
    assert gr.__version__.startswith("5"), f"need gradio 5.x for this smoke, got {gr.__version__}"
    from gradio_client import Client
    medians, metros = _fixtures()
    app_module._DATA["medians"] = medians   # seed cache: no token/network needed
    app_module._DATA["metros"] = metros
    app_module._DATA["error"] = None

    demo = app_module.build_demo()
    demo.launch(prevent_thread_lock=True, server_port=7899, quiet=True)
    try:
        client = Client("http://127.0.0.1:7899")
        cat = client.predict(api_name="/catalog")
        assert any(p["code"] == "45378" for p in cat["procedures"])
        # named-kwarg payload with all params present — must NOT raise "missing parameter"
        res = client.predict(procedure_code="45378", metro="Denver", api_name="/lookup")
        assert res["status"] == "ok" and res["headline"]
        # /ask shares the identical registration path; not invoked (no model in test env)
    finally:
        demo.close()


# ---------------------------------------------------------------------------
# Part K — M3: build_ask graceful LLM failure
# ---------------------------------------------------------------------------

def test_build_ask_llm_failure_is_graceful():
    medians, metros = _fixtures()
    def boom(system, user):
        raise RuntimeError("model down")
    res = app_module.build_ask(medians, metros, "cheapest colonoscopy in denver", CFG, boom)
    assert res["status"] == "unavailable"
    assert res["rows"] == [] and res["recommendation"] is None


# ---------------------------------------------------------------------------
# Part K — M1: suppressed-row exclusion from grounding set
# ---------------------------------------------------------------------------

def test_grounding_excludes_suppressed_row_prices():
    # Tiny Town ($50, suppressed) must NOT legitimize an ungrounded "$50" claim.
    r = _denver_ranking()
    gen = _fake_gen("Skip the wait and pay just $50 somewhere.")  # 50 = a suppressed row's price
    out = llm.write_recommendation(r, CFG, fallback_text="FALLBACK", generate=gen)
    assert out == "FALLBACK"