tests/test_mckesson_regression.py

"""Regression test for the McKesson custom-prompt failure (2026-07).

What actually happened in production (pre-rewrite): a custom prompt that
specified six exact tables (DIM_CUSTOMER_ACCOUNT ... FACT_PRODUCT_USAGE)
was keyword-routed by scenario/selector.py into the healthcare_operations
family, producing FACILITIES / SERVICE_LINES / PAYERS / REGIONS /
HEALTHCARE_OPERATIONS — a generic hospital demo for a company that SELLS
data products.

This test proves the rewrite handles it three ways:
  1. Directives are extracted from the user's actual prompt text.
  2. The exact bad output the old app produced is REJECTED with
     USER DIRECTIVE VIOLATED problems (which the pipeline treats as
     fatal — it can never silently ship).
  3. A blueprint that honors the prompt (6 named tables, 2 facts over
     shared dimensions) validates, generates, and proves its insights.

Run: PYTHONPATH=. python tests/test_mckesson_regression.py
"""

from __future__ import annotations

import sys

from demoprep_app.dataset.engine import BlueprintEngine
from demoprep_app.dataset.validator import validate_bundle
from demoprep_app.scenario.blueprint import DemoBlueprint
from demoprep_app.scenario.directives import directive_problems, extract_directives

# The user's original custom prompt, abbreviated but with every schema-
# bearing phrase intact.
MCKESSON_PROMPT = """Create a McKesson healthcare data products demo for AI-ready analytics
and embedded white-label customer analytics.

McKesson sources, cleanses, links, and sells healthcare data products to pharmaceutical
firms, professional services firms, and large health systems.

Keep the dataset focused and reliable. Create exactly 6 tables:

DIM_CUSTOMER_ACCOUNT
DIM_PROVIDER
DIM_PATIENT_COHORT
DIM_DATA_PRODUCT
FACT_PATIENT_JOURNEY
FACT_PRODUCT_USAGE

Important demo stories:
- Identify patient cohorts with delayed therapy starts and adherence gaps.
- Show provider records with incomplete, stale, or duplicate reference data.
- Show customers using embedded analytics and natural language questioning.
Use synthetic data only. No real patient names and no PHI."""

# The follow-up correction prompt from the transcript.
MCKESSON_CORRECTION = """Create a McKesson data commercialization demo, not a hospital operations demo.

Do not create facilities, service lines, payers, regions, or healthcare operations tables.

McKesson sells cleansed healthcare data products to pharmaceutical companies, consulting
firms, and large health organizations."""

# The exact wrong output the old pipeline produced.
BAD_BLUEPRINT = {
    "company_name": "McKesson",
    "date": {"grain": "day", "months_of_history": 24},
    "dimensions": [
        {"name": "FACILITIES", "values": ["Facility A", "Facility B", "Facility C", "Facility D"]},
        {"name": "SERVICE_LINES", "values": ["Cardiology", "Oncology", "Orthopedics", "Emergency"]},
        {"name": "PAYERS", "values": ["Medicare", "Medicaid", "Commercial", "Self-Pay"]},
        {"name": "REGIONS", "values": ["Northeast", "South", "Midwest", "West"]},
    ],
    "facts": [{
        "name": "HEALTHCARE_OPERATIONS",
        "grain": "facility-service-line-day",
        "dimension_names": ["FACILITIES", "SERVICE_LINES", "PAYERS", "REGIONS"],
        "measures": [{"name": "PATIENT_VISITS", "kind": "base", "distribution": "lognormal", "params": {"mean": 120}, "fmt": "int"}],
    }],
    "insights": [], "dashboard_questions": [],
}

GOOD_BLUEPRINT = {
    "company_name": "McKesson",
    "company_url": "https://www.mckesson.com",
    "use_case": "Healthcare data products demo — AI-ready analytics and embedded customer analytics",
    "business_domain": "healthcare_data_products",
    "demo_audience": "Chief Data Officer and data product GMs",
    "business_problem": "Show patient journey visibility, provider reference quality, embedded/NLQ adoption, and expansion opportunity across data products and customers.",
    "date": {"grain": "month", "months_of_history": 24, "include_date_dimension": False},
    "dimensions": [
        {
            "name": "DIM_CUSTOMER_ACCOUNT",
            "attribute_columns": ["CUSTOMER_TYPE", "CONTRACT_TIER", "REGION"],
            "values": [
                {"name": "Meridian Pharma Insights", "attributes": {"CUSTOMER_TYPE": "Pharmaceutical", "CONTRACT_TIER": "Enterprise", "REGION": "Northeast"}, "performance": 1.5},
                {"name": "Cobalt Health Advisory", "attributes": {"CUSTOMER_TYPE": "Professional Services", "CONTRACT_TIER": "Growth", "REGION": "Midwest"}},
                {"name": "Lakeshore Integrated Health", "attributes": {"CUSTOMER_TYPE": "Health System", "CONTRACT_TIER": "Enterprise", "REGION": "Midwest"}},
                {"name": "Praxis Market Access", "attributes": {"CUSTOMER_TYPE": "Consulting", "CONTRACT_TIER": "Starter", "REGION": "West"}, "performance": 0.6, "mix_weight": 0.7},
                {"name": "Helix Therapeutics Analytics", "attributes": {"CUSTOMER_TYPE": "Pharmaceutical", "CONTRACT_TIER": "Growth", "REGION": "South"}},
            ],
        },
        {
            "name": "DIM_PROVIDER",
            "attribute_columns": ["SPECIALTY", "CARE_SETTING", "DUPLICATE_RISK_FLAG"],
            "values": [
                {"name": "Dr. Elena Vasquez", "attributes": {"SPECIALTY": "Oncology", "CARE_SETTING": "Academic Medical Center", "DUPLICATE_RISK_FLAG": "N"}},
                {"name": "Dr. Marcus Chen", "attributes": {"SPECIALTY": "Cardiology", "CARE_SETTING": "Community Hospital", "DUPLICATE_RISK_FLAG": "N"}},
                {"name": "Dr. Priya Raman", "attributes": {"SPECIALTY": "Endocrinology", "CARE_SETTING": "Ambulatory Clinic", "DUPLICATE_RISK_FLAG": "Y"}, "performance": 0.8},
                {"name": "Dr. James Okafor", "attributes": {"SPECIALTY": "Neurology", "CARE_SETTING": "Academic Medical Center", "DUPLICATE_RISK_FLAG": "N"}},
                {"name": "Dr. Sarah Lindqvist", "attributes": {"SPECIALTY": "Rheumatology", "CARE_SETTING": "Specialty Practice", "DUPLICATE_RISK_FLAG": "Y"}, "performance": 0.75},
            ],
        },
        {
            "name": "DIM_PATIENT_COHORT",
            "attribute_columns": ["DISEASE_AREA", "THERAPY_CLASS", "JOURNEY_STAGE"],
            "values": [
                {"name": "T2D GLP-1 Naive Starters", "attributes": {"DISEASE_AREA": "Type 2 Diabetes", "THERAPY_CLASS": "GLP-1", "JOURNEY_STAGE": "Therapy Start"}, "mix_weight": 1.5},
                {"name": "HFrEF Quad Therapy Gap", "attributes": {"DISEASE_AREA": "Heart Failure", "THERAPY_CLASS": "SGLT2", "JOURNEY_STAGE": "Adherence Gap"}, "performance": 1.3},
                {"name": "NSCLC Biomarker Tested", "attributes": {"DISEASE_AREA": "Oncology", "THERAPY_CLASS": "Targeted Therapy", "JOURNEY_STAGE": "Diagnosis"}},
                {"name": "RA Biologic Switchers", "attributes": {"DISEASE_AREA": "Rheumatoid Arthritis", "THERAPY_CLASS": "Biologic", "JOURNEY_STAGE": "Therapy Switch"}},
            ],
        },
        {
            "name": "DIM_DATA_PRODUCT",
            "attribute_columns": ["PRODUCT_FAMILY", "INCLUDES_EHR_FLAG"],
            "values": [
                {"name": "Open Claims Feed", "attributes": {"PRODUCT_FAMILY": "Patient Data", "INCLUDES_EHR_FLAG": "N"}, "mix_weight": 1.6},
                {"name": "Closed Claims Complete", "attributes": {"PRODUCT_FAMILY": "Patient Data", "INCLUDES_EHR_FLAG": "N"}, "performance": 1.3},
                {"name": "EHR Journey Linked", "attributes": {"PRODUCT_FAMILY": "Patient Data", "INCLUDES_EHR_FLAG": "Y"}},
                {"name": "Provider Reference Master", "attributes": {"PRODUCT_FAMILY": "Reference Data", "INCLUDES_EHR_FLAG": "N"}, "mix_weight": 1.2},
            ],
        },
    ],
    "facts": [
        {
            "name": "FACT_PATIENT_JOURNEY",
            "grain": "one row per customer per cohort per provider per data product per month",
            "dimension_names": ["DIM_CUSTOMER_ACCOUNT", "DIM_PATIENT_COHORT", "DIM_PROVIDER", "DIM_DATA_PRODUCT"],
            "date_column": "MONTH_DATE",
            "measures": [
                {"name": "PATIENTS_OBSERVED", "kind": "base", "distribution": "lognormal", "params": {"mean": 850, "sigma": 0.6}, "fmt": "int"},
                {"name": "THERAPY_STARTS", "kind": "base", "distribution": "lognormal", "params": {"mean": 60, "sigma": 0.7}, "fmt": "int"},
                {"name": "ADHERENCE_GAP_PATIENTS", "kind": "base", "distribution": "lognormal", "params": {"mean": 95, "sigma": 0.8}, "fmt": "int"},
                {"name": "ESTIMATED_OPPORTUNITY_USD", "kind": "base", "distribution": "lognormal", "params": {"mean": 68000, "sigma": 0.9}, "fmt": "currency"},
            ],
        },
        {
            "name": "FACT_PRODUCT_USAGE",
            "grain": "one row per customer per data product per month",
            "dimension_names": ["DIM_CUSTOMER_ACCOUNT", "DIM_DATA_PRODUCT"],
            "date_column": "MONTH_DATE",
            "measures": [
                {"name": "QUERIES_RUN", "kind": "base", "distribution": "lognormal", "params": {"mean": 4200, "sigma": 0.7}, "fmt": "int"},
                {"name": "NL_QUESTIONS", "kind": "base", "distribution": "lognormal", "params": {"mean": 380, "sigma": 0.8}, "fmt": "int"},
                {"name": "EMBEDDED_SESSIONS", "kind": "base", "distribution": "lognormal", "params": {"mean": 240, "sigma": 0.7}, "fmt": "int"},
                {"name": "DATA_QUALITY_ISSUES", "kind": "base", "distribution": "lognormal", "params": {"mean": 6, "sigma": 0.9}, "fmt": "int"},
                {"name": "EXPANSION_OPPORTUNITY_USD", "kind": "base", "distribution": "lognormal", "params": {"mean": 42000, "sigma": 0.9}, "fmt": "currency"},
            ],
        },
    ],
    "seasonality": {
        "monthly": {"1": 1.05, "2": 1.0, "3": 1.1, "4": 1.0, "5": 0.98, "6": 0.95,
                    "7": 0.9, "8": 0.92, "9": 1.05, "10": 1.1, "11": 1.05, "12": 0.95},
        "trend_pct_per_year": 15,
        "narrative": ["Steady usage growth as accounts onboard", "Summer analytics lull", "Q4 planning-season lift"],
    },
    "insights": [
        {
            "id": "hf_adherence_gap",
            "headline": "Heart-failure cohorts show the widest adherence gaps",
            "insight_type": "segment_outlier",
            "dimension": "DIM_PATIENT_COHORT", "dimension_value": "HFrEF Quad Therapy Gap",
            "measure": "ADHERENCE_GAP_PATIENTS", "magnitude": 1.5,
            "spotter_question": "Show adherence gap patients by patient cohort",
            "expected_finding": "HFrEF Quad Therapy Gap is the tallest bar by a wide margin",
        },
        {
            "id": "nlq_takeoff",
            "headline": "Natural-language questioning surged after the embedded analytics relaunch",
            "insight_type": "trend_break",
            "dimension": "DIM_CUSTOMER_ACCOUNT", "dimension_value": "Meridian Pharma Insights",
            "measure": "NL_QUESTIONS", "magnitude": 1.6, "window_start": "2026-01-01",
            "spotter_question": "Show natural language questions by customer account by month",
            "expected_finding": "Meridian's NLQ volume breaks upward from January 2026",
        },
        {
            "id": "reference_quality_drag",
            "headline": "Providers flagged for duplicate risk drag reference-data value",
            "insight_type": "segment_laggard",
            "dimension": "DIM_PROVIDER", "dimension_value": "Dr. Sarah Lindqvist",
            "measure": "PATIENTS_OBSERVED", "magnitude": 0.7,
            "spotter_question": "Show patients observed by provider",
            "expected_finding": "Duplicate-flagged providers sit at the bottom of the ranking",
        },
    ],
    "dashboard_questions": [
        "Which patient cohorts have delayed therapy starts and adherence gaps?",
        "Where do open claims, closed claims, and EHR data improve journey visibility?",
        "Which customers are adopting embedded analytics and natural language questioning?",
        "Where is expansion opportunity concentrated by customer, product, and disease area?",
        "Which provider records have data quality issues?",
    ],
}


def main() -> int:
    failures = []

    # ------------------------------------------------------------------
    # 1. Directive extraction from the user's ACTUAL prompts
    # ------------------------------------------------------------------
    d1 = extract_directives(MCKESSON_PROMPT)
    print(f"[1a] original prompt -> required tables: {d1.required_tables}")
    expected = {"DIM_CUSTOMER_ACCOUNT", "DIM_PROVIDER", "DIM_PATIENT_COHORT",
                "DIM_DATA_PRODUCT", "FACT_PATIENT_JOURNEY", "FACT_PRODUCT_USAGE"}
    if set(d1.required_tables) != expected:
        failures.append(f"expected {expected}, got {set(d1.required_tables)}")

    d2 = extract_directives(MCKESSON_CORRECTION)
    print(f"[1b] correction prompt -> prohibited: {d2.prohibited_terms}")
    for term in ("FACILITIES", "SERVICE_LINES", "PAYERS", "HEALTHCARE_OPERATIONS"):
        if not any(term in p or p in term for p in d2.prohibited_terms):
            failures.append(f"prohibition '{term}' not extracted from correction prompt")

    # ------------------------------------------------------------------
    # 2. The exact bad output the old app produced must be REJECTED
    # ------------------------------------------------------------------
    bad = DemoBlueprint.from_dict(BAD_BLUEPRINT, seed=1)
    viol_required = directive_problems(bad, d1)   # missing every named table
    viol_prohibited = directive_problems(bad, d2) # contains every prohibited table
    print(f"[2] bad blueprint violations: {len(viol_required)} missing-table, {len(viol_prohibited)} prohibited-table")
    for v in (viol_required + viol_prohibited)[:4]:
        print("     ", v)
    if len(viol_required) != 6:
        failures.append(f"expected 6 missing-table violations, got {len(viol_required)}")
    if len(viol_prohibited) < 4:
        failures.append(f"expected >=4 prohibited-table violations, got {len(viol_prohibited)}")
    if not all("USER DIRECTIVE VIOLATED" in v for v in viol_required + viol_prohibited):
        failures.append("directive violations must carry the structural marker the pipeline aborts on")

    # ------------------------------------------------------------------
    # 3. A blueprint honoring the prompt validates, generates 2 facts,
    #    and proves its insights
    # ------------------------------------------------------------------
    good = DemoBlueprint.from_dict(GOOD_BLUEPRINT, seed=77)
    problems = good.problems() + directive_problems(good, d1) + directive_problems(good, d2)
    print(f"[3a] good blueprint problems: {problems or 'clean'}")
    if problems:
        failures.append(f"good blueprint should be clean: {problems}")

    bundle = BlueprintEngine().generate(good, row_count=10000)
    names = [t.name for t in bundle.tables]
    print(f"[3b] generated tables: {names}")
    if set(names) != {"DIM_CUSTOMER_ACCOUNT", "DIM_PROVIDER", "DIM_PATIENT_COHORT",
                      "DIM_DATA_PRODUCT", "FACT_PATIENT_JOURNEY", "FACT_PRODUCT_USAGE"}:
        failures.append("generated tables do not match the user's 6-table spec")

    journey = bundle.table_map()["FACT_PATIENT_JOURNEY"]
    usage = bundle.table_map()["FACT_PRODUCT_USAGE"]
    print(f"[3c] fact rows: journey={len(journey.rows)}, usage={len(usage.rows)}; "
          f"journey date col: {'MONTH_DATE' in journey.rows[0]}")
    if "MONTH_DATE" not in journey.rows[0] or "MONTH_DATE" not in usage.rows[0]:
        failures.append("custom date_column MONTH_DATE not honored")
    if f"DIM_PROVIDER_KEY" in usage.rows[0]:
        failures.append("FACT_PRODUCT_USAGE should not reference DIM_PROVIDER")

    report = validate_bundle(good, bundle)
    print("[3d]\n" + report.summary())
    if not report.passed:
        failures.append("validation failed:\n" + report.summary())

    if failures:
        print("\nFAILURES:")
        for f in failures:
            print(" -", f)
        return 1
    print("\nMCKESSON REGRESSION: ALL CHECKS PASSED")
    return 0


if __name__ == "__main__":
    sys.exit(main())