test_data.py

"""Unit tests for Hub metrics discovery helpers (no network)."""

from __future__ import annotations

import pandas as pd

from data import (
    condition_family,
    condition_label,
    condition_order_index,
    disagreement_by_slice,
    disagreement_summary,
    encoder_comparator_frame,
    generative_ladder_frame,
    slice_accuracy_frame,
)


def test_rfdt_labels_and_family() -> None:
    assert condition_label("rfdt_base") == "RFDT Base"
    assert condition_label("rfdt_lora") == "RFDT-LoRA"
    assert condition_family("rfdt_base") == "rfdt"
    assert condition_family("rfdt_lora") == "rfdt"
    assert condition_family("lora") == "learned"
    assert condition_family("base") == "frozen"


def test_encoder_labels_and_family() -> None:
    assert condition_label("encoder_ft") == "Encoder-FT"
    assert condition_family("encoder_ft") == "encoder"


def test_unknown_condition_gets_fallback_label() -> None:
    assert condition_label("future_method_v2") == "Future Method V2"
    assert condition_family("future_method_v2") == "other"


def test_rfdt_sorts_after_cartridge() -> None:
    assert condition_order_index("cartridge") < condition_order_index("rfdt_base")
    assert condition_order_index("rfdt_base") < condition_order_index("rfdt_lora")
    assert condition_order_index("rfdt_lora") < condition_order_index("encoder_ft")


def test_generative_ladder_drops_encoder_rows() -> None:
    df = pd.DataFrame(
        [
            {"model_key": "qwen3-8b", "condition": "lora", "family": "learned"},
            {
                "model_key": "modernbert-base",
                "condition": "encoder_ft",
                "family": "encoder",
            },
        ]
    )
    out = generative_ladder_frame(df)
    assert list(out["model_key"]) == ["qwen3-8b"]


def test_encoder_comparator_keeps_rfdt_and_modernbert() -> None:
    df = pd.DataFrame(
        [
            {
                "model_key": "qwen3-8b",
                "condition": "rfdt_lora",
                "family": "rfdt",
            },
            {
                "model_key": "qwen3-8b",
                "condition": "lora",
                "family": "learned",
            },
            {
                "model_key": "modernbert-base",
                "condition": "encoder_ft",
                "family": "encoder",
            },
        ]
    )
    out = encoder_comparator_frame(df)
    assert set(out["condition"]) == {"rfdt_lora", "encoder_ft"}
    assert "lora" not in set(out["condition"])


def test_slice_accuracy_frame_long_form() -> None:
    df = pd.DataFrame(
        [
            {
                "model": "Qwen3-8B",
                "model_key": "qwen3-8b",
                "condition": "rfdt_lora",
                "condition_label": "RFDT-LoRA",
                "slice_acc_allowed_base": 1.0,
                "slice_n_allowed_base": 8,
                "slice_acc_denied_edge": 0.5,
                "slice_n_denied_edge": 66,
                "slice_acc_allowed_edge": None,
                "slice_n_allowed_edge": None,
                "slice_acc_denied_base": 0.9,
                "slice_n_denied_base": 12,
            },
            {
                "model": "ModernBERT-base",
                "model_key": "modernbert-base",
                "condition": "encoder_ft",
                "condition_label": "Encoder-FT",
                "slice_acc_allowed_base": 1.0,
                "slice_n_allowed_base": 8,
                "slice_acc_denied_edge": 0.4,
                "slice_n_denied_edge": 66,
                "slice_acc_allowed_edge": 0.9,
                "slice_n_allowed_edge": 43,
                "slice_acc_denied_base": 0.0,
                "slice_n_denied_base": 12,
            },
        ]
    )
    long = slice_accuracy_frame(df)
    assert not long.empty
    assert set(long["query_type"]) <= {
        "allowed_base",
        "allowed_edge",
        "denied_base",
        "denied_edge",
    }
    assert "Qwen3-8B · RFDT-LoRA" in set(long["method"])
    assert long.loc[
        (long["method"].str.contains("ModernBERT"))
        & (long["query_type"] == "denied_edge"),
        "accuracy",
    ].iloc[0] == 0.4


def test_disagreement_summary_rates() -> None:
    disagree = pd.DataFrame(
        [
            {
                "model": "Qwen3-8B",
                "model_key": "qwen3-8b",
                "query_type": "denied_edge",
                "outcome": "rfdt_only",
            },
            {
                "model": "Qwen3-8B",
                "model_key": "qwen3-8b",
                "query_type": "denied_edge",
                "outcome": "encoder_only",
            },
            {
                "model": "Qwen3-8B",
                "model_key": "qwen3-8b",
                "query_type": "allowed_edge",
                "outcome": "both_correct",
            },
            {
                "model": "Qwen3-8B",
                "model_key": "qwen3-8b",
                "query_type": "allowed_edge",
                "outcome": "both_correct",
            },
        ]
    )
    summary = disagreement_summary(disagree)
    assert abs(summary.loc[summary["outcome"] == "both_correct", "rate"].iloc[0] - 0.5) < 1e-9
    by_slice = disagreement_by_slice(disagree)
    denied = by_slice[by_slice["query_type"] == "denied_edge"].iloc[0]
    assert abs(denied["net_rfdt_edge"] - 0.0) < 1e-9
    assert abs(denied["disagree_rate"] - 1.0) < 1e-9