tests/test_tabular_summary.py

from __future__ import annotations

import pandas as pd
import pytest

from sepsis_mcp.tabular_summary import build_summary_tables, main


def test_build_summary_tables_merges_run_level_and_subgroup_metrics(tmp_path) -> None:
    sweep_dir = tmp_path / "sweep"
    runs_dir = sweep_dir / "runs" / "cal64__xgboost__alpha0p1__none"
    runs_dir.mkdir(parents=True)

    pd.DataFrame(
        [
            {
                "run_id": "cal64__xgboost__alpha0p1__none",
                "experiment": "A_to_B",
                "method": "standard",
                "model_type": "xgboost",
                "calibration_patients": 64,
                "alpha": 0.1,
                "mask_strategy": "none",
                "auroc": 0.76,
                "auprc": 0.05,
                "empirical_coverage": 0.92,
                "average_set_size": 0.94,
                "abstention_rate": 0.0,
                "max_group_coverage_gap": 0.04,
                "target_coverage": 0.9,
            },
            {
                "run_id": "cal64__xgboost__alpha0p1__none",
                "experiment": "A_to_B",
                "method": "missingness_aware",
                "model_type": "xgboost",
                "calibration_patients": 64,
                "alpha": 0.1,
                "mask_strategy": "none",
                "auroc": 0.76,
                "auprc": 0.05,
                "empirical_coverage": 0.89,
                "average_set_size": 0.97,
                "abstention_rate": 0.0,
                "max_group_coverage_gap": 0.06,
                "target_coverage": 0.9,
            },
        ]
    ).to_csv(sweep_dir / "sweep_results.csv", index=False)

    pd.DataFrame(
        [
            {
                "experiment": "A_to_B",
                "method": "standard",
                "group": 0,
                "group_label": "low",
                "count": 20,
                "coverage": 0.95,
            },
            {
                "experiment": "A_to_B",
                "method": "standard",
                "group": 2,
                "group_label": "high",
                "count": 12,
                "coverage": 0.87,
            },
            {
                "experiment": "A_to_B",
                "method": "missingness_aware",
                "group": 0,
                "group_label": "low",
                "count": 20,
                "coverage": 0.91,
            },
            {
                "experiment": "A_to_B",
                "method": "missingness_aware",
                "group": 2,
                "group_label": "high",
                "count": 12,
                "coverage": 0.9,
            },
        ]
    ).to_csv(runs_dir / "subgroup_coverage.csv", index=False)

    overall_summary, subgroup_summary = build_summary_tables(sweep_dir)

    assert set(overall_summary["method"]) == {"standard", "missingness_aware"}
    assert {"coverage_gap", "absolute_coverage_gap"} <= set(overall_summary.columns)
    assert overall_summary.loc[overall_summary["method"] == "standard", "coverage_gap"].item() == pytest.approx(
        0.02
    )

    assert set(subgroup_summary["group_label"]) == {"low", "high"}
    assert {"model_type", "calibration_patients", "auroc", "coverage_gap", "is_high_missing_group"} <= set(
        subgroup_summary.columns
    )
    high_row = subgroup_summary[
        (subgroup_summary["method"] == "standard") & (subgroup_summary["group_label"] == "high")
    ].iloc[0]
    assert high_row["coverage_gap"] == pytest.approx(-0.03)
    assert bool(high_row["is_high_missing_group"]) is True


def test_tabular_summary_main_writes_overall_and_subgroup_csvs(tmp_path) -> None:
    sweep_dir = tmp_path / "sweep"
    runs_dir = sweep_dir / "runs" / "cal64__xgboost__alpha0p1__none"
    runs_dir.mkdir(parents=True)

    pd.DataFrame(
        [
            {
                "run_id": "cal64__xgboost__alpha0p1__none",
                "experiment": "A_to_A",
                "method": "standard",
                "model_type": "xgboost",
                "calibration_patients": 64,
                "alpha": 0.1,
                "mask_strategy": "none",
                "auroc": 0.77,
                "auprc": 0.06,
                "empirical_coverage": 0.91,
                "average_set_size": 0.95,
                "abstention_rate": 0.0,
                "max_group_coverage_gap": 0.03,
                "target_coverage": 0.9,
            }
        ]
    ).to_csv(sweep_dir / "sweep_results.csv", index=False)

    pd.DataFrame(
        [
            {
                "experiment": "A_to_A",
                "method": "standard",
                "group": 1,
                "group_label": "medium",
                "count": 10,
                "coverage": 0.92,
            }
        ]
    ).to_csv(runs_dir / "subgroup_coverage.csv", index=False)

    output_dir = tmp_path / "summary"
    main(
        [
            "--sweep-dir",
            str(sweep_dir),
            "--output-dir",
            str(output_dir),
        ]
    )

    overall_path = output_dir / "overall_summary.csv"
    subgroup_path = output_dir / "subgroup_summary.csv"

    assert overall_path.exists()
    assert subgroup_path.exists()
    assert pd.read_csv(overall_path).shape[0] == 1
    assert pd.read_csv(subgroup_path).shape[0] == 1