tests/test_inference.py

"""Inference adapter tests (D-MASK-01..04, D-CAL-08, D-CAL-09)."""

import numpy as np
import pytest

from model.features import CLASSES
from model.inference import MASK_TABLE, apply_mask_and_renormalize


def test_mask_table_completeness() -> None:
    """All 4 NetworkMode values are keys; every value is a frozenset of valid CLASSES."""
    assert set(MASK_TABLE.keys()) == {"enterprise", "captive", "home", "unknown"}
    for mode, applicable in MASK_TABLE.items():
        assert isinstance(applicable, frozenset)
        assert applicable.issubset(set(CLASSES)), (
            f"{mode} contains non-canonical class slug"
        )


def test_mask_table_unknown_all_ten() -> None:
    """D-MASK-03: unknown mode = all 10 CLASSES enabled."""
    assert MASK_TABLE["unknown"] == frozenset(CLASSES)


def test_mask_table_enterprise() -> None:
    """D-MASK-02: enterprise = 8 classes; captive_portal_expiry and isp_upstream_fail excluded."""
    ent = MASK_TABLE["enterprise"]
    assert len(ent) == 8
    assert "captive_portal_expiry" not in ent
    assert "isp_upstream_fail" not in ent
    assert "auth_8021x_eap_fail" in ent
    assert "ap_roam_rekey_fail" in ent
    assert "radius_timeout" in ent


def test_mask_table_captive() -> None:
    """D-MASK-02: captive = 5 classes."""
    assert MASK_TABLE["captive"] == frozenset({
        "captive_portal_expiry", "dns_resolver_fail", "isp_upstream_fail",
        "dhcp_lease_churn", "mac_randomization_reject",
    })


def test_mask_table_home() -> None:
    """D-MASK-02: home = 5 classes."""
    assert MASK_TABLE["home"] == frozenset({
        "dhcp_lease_churn", "dns_resolver_fail", "driver_power_save_wake",
        "rf_sticky_client", "isp_upstream_fail",
    })


def _uniform_probs() -> np.ndarray:
    return np.full(10, 1.0 / 10)


@pytest.mark.parametrize("mode", ["enterprise", "captive", "home", "unknown"])
def test_mask_renormalize_sums_to_one(mode: str) -> None:
    """CLASS-04: every mode produces probs that sum to 1.0 within 1e-9."""
    out = apply_mask_and_renormalize(_uniform_probs(), mode)
    assert abs(out.sum() - 1.0) < 1e-9


def test_mask_table_enterprise_zeros_excluded() -> None:
    """CLASS-04: enterprise mask zeros captive_portal_expiry + isp_upstream_fail."""
    out = apply_mask_and_renormalize(_uniform_probs(), "enterprise")
    cap_idx = CLASSES.index("captive_portal_expiry")
    isp_idx = CLASSES.index("isp_upstream_fail")
    assert out[cap_idx] == 0.0
    assert out[isp_idx] == 0.0


def test_mask_unknown_passthrough() -> None:
    """D-MASK-03: unknown mode is no-op for already-normalized probs."""
    u = _uniform_probs()
    out = apply_mask_and_renormalize(u, "unknown")
    np.testing.assert_allclose(out, u, atol=1e-12)


def test_mask_renormalize_preserves_relative_order() -> None:
    """Within-mask classes keep their relative ranking after renormalization."""
    # Crafted probs: enterprise classes get rising values; non-enterprise get noise.
    probs = np.array([0.05, 0.10, 0.15, 0.99, 0.20, 0.25, 0.08, 0.30, 0.35, 0.99])
    # CLASSES order: [auth_8021x..., ap_roam..., radius..., captive_portal_expiry,
    #                 mac_rand..., dhcp..., dns..., driver..., rf_sticky..., isp_upstream]
    out = apply_mask_and_renormalize(probs, "enterprise")
    # captive_portal_expiry (idx 3) and isp_upstream_fail (idx 9) should be 0
    assert out[3] == 0.0
    assert out[9] == 0.0
    # rf_sticky_client (idx 8) was highest among enterprise classes; should remain highest
    assert int(np.argmax(out)) == 8


# CLASS-03: top_k full ranking — guarded behind data/train.parquet existence
# because we need a trained classifier to call predict_verdict.
@pytest.mark.skipif(
    not __import__("pathlib").Path("data/train.parquet").exists(),
    reason="data/train.parquet not generated (run `make synth` first)",
)
def test_top_k_full_ranking(tmp_path) -> None:  # type: ignore[no-untyped-def]
    """CLASS-03: predict_verdict returns Verdict with len(top_k) == 10 ordered by prob."""
    from pathlib import Path

    import joblib

    from model.features import load_split
    from model.inference import predict_verdict
    from model.train_classifier import train_calibrated_classifier

    X, y, _ = load_split(Path("data/train.parquet"))
    # Fast smoke: 50 rows per class
    idx_list: list[int] = []
    for c in range(10):
        idx_list.extend(np.where(y == c)[0][:50].tolist())
    idx_arr = np.array(idx_list)
    clf = train_calibrated_classifier(
        X[idx_arr], y[idx_arr], classifier_seed=42, cv_seed=43
    )
    joblib_path = tmp_path / "classifier.joblib"
    joblib.dump(clf, joblib_path, compress=3)

    # Build a synthetic frame list (N=3 frames, network_mode=enterprise)
    frames: list[dict] = []
    for _ in range(3):
        row = {
            "rssi_dbm": -65.0,
            "ping_continuity": {
                "window_ms": 30000, "avg_rtt_ms": 25.0,
                "packet_loss_pct": 0.0, "jitter_ms": 5.0,
            },
            "latency_jitter_ms": 5.0, "dns_resolution_ms": 12.0,
            "per_packet_retry_count": 1, "beacon_rssi_dbm": -60.0,
            "neighbor_ap_count_5ghz": 3,
            "os": "windows", "network_mode": "enterprise",
            "dhcp_event_class": "none", "auth_event_class": "none",
            "mac_randomization_state": "off", "driver_state": "normal",
            "captive_portal_detected": False, "bssid_mode": "raw",
            "window_ms": 30000, "channel": 36, "rts_cts_rate": 0.05,
        }
        frames.append(row)

    verdict = predict_verdict(str(joblib_path), frames)

    # CLASS-03: full ranking, K=10
    assert len(verdict.top_k) == 10
    # Ordered by descending prob
    probs_in_order = [p for _, p in verdict.top_k]
    assert probs_in_order == sorted(probs_in_order, reverse=True)
    # top_class / confidence consistent with top_k[0]
    assert verdict.top_class == verdict.top_k[0][0]
    assert verdict.confidence == verdict.top_k[0][1]
    # Mode=enterprise: captive_portal_expiry and isp_upstream_fail must have prob 0.0
    d = dict(verdict.top_k)
    assert d["captive_portal_expiry"] == 0.0
    assert d["isp_upstream_fail"] == 0.0
    # Sums to 1.0 (post-mask renormalization)
    assert abs(sum(p for _, p in verdict.top_k) - 1.0) < 1e-9