tests/test_pr2_analytics.py

import pandas as pd
import plotly.graph_objects as go

from src.analytics import (
    BENCHMARK_CATALOG,
    benchmark_category,
    benchmarks_for_category,
    cross_benchmark_ranking_df,
    enrich_analysis_df,
    filter_category,
    matrix_df,
    ranking_df,
)
from src.charts import create_coverage_matrix_plot, create_matrix_plot, create_tradeoff_plot


def frame():
    return pd.DataFrame(
        [
            {
                "Benchmark": "SWE-Bench Verified", "Model": "a", "Harness": "h", "Run Label": "a / h",
                "Category": "FOSS", "Score": .8, "Score (%)": 80, "Tasks": 10, "Errors": 1,
                "Input Tokens Per Task": 50, "Cache Tokens Per Task": 10, "Output Tokens Per Task": 20,
                "Total Tokens Per Task": 80, "Tokens Per Solved Task": 100, "Cost Per Task": .2,
                "Total Time Per Task": 10, "Agent Time Per Task": 8, "Token Data Available": True,
            },
            {
                "Benchmark": "SWE-Bench Verified", "Model": "b", "Harness": "h", "Run Label": "b / h",
                "Category": "FOSS", "Score": .4, "Score (%)": 40, "Tasks": 10, "Errors": 2,
                "Input Tokens Per Task": 90, "Cache Tokens Per Task": None, "Output Tokens Per Task": 30,
                "Total Tokens Per Task": 120, "Tokens Per Solved Task": 300, "Cost Per Task": .1,
                "Total Time Per Task": 20, "Agent Time Per Task": 15, "Token Data Available": True,
            },
            {
                "Benchmark": "Terminal Bench 2.0", "Model": "a", "Harness": "h", "Run Label": "a / h",
                "Category": "FOSS", "Score": .2, "Score (%)": 20, "Tasks": 10, "Errors": 0,
                "Input Tokens Per Task": 30, "Cache Tokens Per Task": 5, "Output Tokens Per Task": 10,
                "Total Tokens Per Task": 45, "Tokens Per Solved Task": 225, "Cost Per Task": .3,
                "Total Time Per Task": 30, "Agent Time Per Task": 25, "Token Data Available": True,
            },
            {
                "Benchmark": "Terminal Bench 2.0", "Model": "c", "Harness": "h", "Run Label": "c / h",
                "Category": "FOSS", "Score": .9, "Score (%)": 90, "Tasks": 10, "Errors": None,
                "Input Tokens Per Task": None, "Cache Tokens Per Task": None, "Output Tokens Per Task": None,
                "Total Tokens Per Task": None, "Tokens Per Solved Task": None, "Cost Per Task": None,
                "Total Time Per Task": None, "Agent Time Per Task": None, "Token Data Available": False,
            },
            {
                "Benchmark": "New Benchmark", "Model": "z", "Harness": "h", "Run Label": "z / h",
                "Category": "FOSS", "Score": .5, "Score (%)": 50, "Tasks": 10, "Errors": 0,
                "Input Tokens Per Task": 1, "Cache Tokens Per Task": 1, "Output Tokens Per Task": 1,
                "Total Tokens Per Task": 3, "Tokens Per Solved Task": 6, "Cost Per Task": .01,
                "Total Time Per Task": 1, "Agent Time Per Task": 1, "Token Data Available": True,
            },
        ]
    )


def test_benchmark_catalog_and_unknown_fallback():
    assert benchmark_category("SWE-Bench Verified") == "Coding"
    assert benchmark_category("Terminal Bench 2.0") == "Generalist"
    assert benchmark_category("Shellbench") == "Generalist"
    assert benchmark_category("New Benchmark") == "Other"
    assert "SWE-Bench Pro -- Ansible" in BENCHMARK_CATALOG


def test_category_filtering_keeps_unknown_visible_as_other():
    df = frame()
    assert set(benchmarks_for_category(df, "Coding")) == {"SWE-Bench Verified"}
    assert set(filter_category(df, "Generalist")["Benchmark"]) == {"Terminal Bench 2.0"}
    assert set(filter_category(df, "Other")["Benchmark"]) == {"New Benchmark"}


def test_derived_reliability_and_per_success_metrics():
    df = enrich_analysis_df(frame())
    first = df.iloc[0]
    assert first["Execution Error Rate (%)"] == 10
    assert first["Tokens Per Successful Task"] == 100
    assert first["Cost Per Successful Task"] == .25
    assert first["Time Per Successful Task"] == 12.5
    assert pd.isna(df.loc[df["Model"].eq("c"), "Execution Error Rate (%)"]).all()


def test_within_benchmark_percentile_and_rank():
    df = enrich_analysis_df(frame())
    coding = df[df["Benchmark"] == "SWE-Bench Verified"].set_index("Model")
    assert coding.loc["a", "Within-Benchmark Rank"] == 1
    assert coding.loc["b", "Within-Benchmark Rank"] == 2
    assert coding.loc["a", "Within-Benchmark Percentile"] == 100
    assert coding.loc["b", "Within-Benchmark Percentile"] == 0


def test_cross_benchmark_ordering_uses_normalization_and_coverage_threshold():
    df = frame()
    ranked = cross_benchmark_ranking_df(df, minimum_coverage=0.5)
    a = ranked[ranked["Model"] == "a"].iloc[0]
    assert a["Benchmarks Covered"] == 2
    assert a["Eligible Benchmarks"] == 3
    assert a["Normalized Performance"] == 50
    strict = cross_benchmark_ranking_df(df, minimum_coverage=1.0)
    assert strict.empty


def test_token_ranking_excludes_missing_and_orders_lower_first():
    ranked = ranking_df(frame(), "Total tokens", benchmark="Terminal Bench 2.0")
    assert ranked["Model"].tolist() == ["a"]


def test_score_rank_and_metric_matrices_preserve_missing_cells():
    df = frame()
    score = matrix_df(df, "Score")
    rank = matrix_df(df, "Within-benchmark rank")
    cost = matrix_df(df, "Cost")
    assert pd.isna(score.loc["b / h", "Terminal Bench 2.0"])
    assert rank.loc["a / h", "SWE-Bench Verified"] == 1
    assert pd.isna(cost.loc["c / h", "Terminal Bench 2.0"])


def test_coverage_matrix_uses_missing_not_zero_score():
    matrix = matrix_df(frame(), "Coverage")
    assert matrix.loc["a / h", "SWE-Bench Verified"] == 1
    assert pd.isna(matrix.loc["b / h", "Terminal Bench 2.0"])
    figure = create_coverage_matrix_plot(matrix)
    assert isinstance(figure, go.Figure)
    assert "Available" in figure.data[0].text[0] or "Missing" in figure.data[0].text[0]


def test_generic_tradeoff_and_matrix_charts_construct():
    df = enrich_analysis_df(frame())
    figure = create_tradeoff_plot(
        df[df["Benchmark"] == "SWE-Bench Verified"],
        "Total Tokens Per Task", "Score (%)",
        "Total tokens per task", "Score (%)",
        show_pareto_frontier=True,
    )
    matrix_figure = create_matrix_plot(matrix_df(df, "Score"), "Score matrix", "Score (%)")
    assert isinstance(figure, go.Figure)
    assert isinstance(matrix_figure, go.Figure)
    assert any(trace.name == "Pareto frontier" for trace in figure.data)


def test_tradeoff_pareto_supports_lower_is_better_on_both_axes():
    df = pd.DataFrame(
        {
            "Model": ["a", "b", "c"],
            "Harness": ["h", "h", "h"],
            "Benchmark": ["bench", "bench", "bench"],
            "Total Tokens Per Task": [100, 200, 300],
            "Cost Per Task": [0.3, 0.2, 0.4],
        }
    )
    figure = create_tradeoff_plot(
        df,
        "Total Tokens Per Task",
        "Cost Per Task",
        "Tokens",
        "Cost",
        show_pareto_frontier=True,
        lower_x_is_better=True,
        higher_y_is_better=False,
    )
    frontier = next(trace for trace in figure.data if trace.name == "Pareto frontier")
    assert list(frontier.x) == [100, 200]
    assert list(frontier.y) == [0.3, 0.2]


def test_matrix_can_color_by_percentile_but_display_raw_values():
    color_matrix = pd.DataFrame([[100.0, 0.0]], index=["a / h"], columns=["b1", "b2"] )
    raw_matrix = pd.DataFrame([[82.5, 41.25]], index=["a / h"], columns=["b1", "b2"] )
    figure = create_matrix_plot(
        color_matrix,
        "Score matrix",
        "Within-benchmark percentile",
        display_matrix=raw_matrix,
        display_metric_label="Benchmark score (%)",
    )
    assert list(figure.data[0].z[0]) == [100.0, 0.0]
    assert list(figure.data[0].text[0]) == ["82.5", "41.25"]
    assert "Benchmark score (%)" in figure.data[0].hovertemplate


def test_ranking_plot_order_can_be_value_or_alphabetical():
    from src.charts import create_ranking_plot

    df = pd.DataFrame(
        {
            "Model": ["b", "a", "c"],
            "Harness": ["h", "h", "h"],
            "Benchmark": ["bench", "bench", "bench"],
            "Score (%)": [20, 10, 30],
        }
    )
    largest = create_ranking_plot(
        df, "Score (%)", "Score", True, sort_order="Largest value first"
    )
    lowest = create_ranking_plot(
        df, "Score (%)", "Score", True, sort_order="Lowest value first"
    )
    alpha = create_ranking_plot(
        df, "Score (%)", "Score", True, sort_order="Alphabetical (A–Z)"
    )

    assert list(largest.layout.yaxis.categoryarray) == ["c / h", "b / h", "a / h"]
    assert list(lowest.layout.yaxis.categoryarray) == ["a / h", "b / h", "c / h"]
    assert list(alpha.layout.yaxis.categoryarray) == ["a / h", "b / h", "c / h"]


def test_matrix_height_scales_with_rows_and_keeps_all_y_labels():
    rows = 18
    matrix = pd.DataFrame(
        {"bench-a": range(rows), "bench-b": range(rows)},
        index=[f"model-{i} / harness" for i in range(rows)],
    )

    figure = create_matrix_plot(matrix, "Dense matrix", "Score (%)")

    assert figure.layout.height >= 900
    assert figure.layout.yaxis.tickmode == "array"
    assert len(figure.layout.yaxis.tickvals) == rows
    assert len(figure.layout.yaxis.ticktext) == rows