tests/test_bench.py

"""Tests for the benchmark harness (fast, small-D)."""
from __future__ import annotations

import json
import sys
from pathlib import Path

import numpy as np
import pytest

# make bench/ importable
BENCH_DIR = Path(__file__).resolve().parent.parent / "bench"
sys.path.insert(0, str(BENCH_DIR))

import gen_datasets  # noqa: E402
import run_benchmark as bench  # noqa: E402
from ensemble import Expert  # noqa: E402


@pytest.fixture
def small_expert():
    text = "the mitochondria is the powerhouse of the cell. " * 20
    return Expert.from_text(text, domain="bio", D=2000, seed=0)


class TestDatasets:
    def test_gen_facts_deterministic(self):
        a = gen_datasets.gen_facts()
        b = gen_datasets.gen_facts()
        assert a == b

    def test_gen_qa_splits_disjoint(self):
        train, holdout = gen_datasets.gen_qa_splits()
        train_qs = {q for q, _ in train}
        holdout_qs = {q for q, _ in holdout}
        # holdout questions must NOT appear in train
        assert holdout_qs.isdisjoint(train_qs)

    def test_gen_all(self, tmp_path):
        paths = gen_datasets.generate_all(tmp_path)
        assert (tmp_path / "facts.txt").exists()
        assert (tmp_path / "qa_train.json").exists()
        assert (tmp_path / "qa_holdout.json").exists()
        assert (tmp_path / "prose.txt").exists()


class TestMetrics:
    def test_next_token_accuracy_in_range(self, small_expert):
        acc = bench._next_token_accuracy(small_expert, "the mitochondria is the powerhouse", max_eval=50)
        assert 0.0 <= acc <= 1.0

    def test_next_token_accuracy_repetitive_high(self, small_expert):
        # on the exact training text, accuracy should be decent (memorized)
        text = "the mitochondria is the powerhouse of the cell. "
        acc = bench._next_token_accuracy(small_expert, text * 5, max_eval=80)
        assert acc > 0.3  # repeated text is well learned

    def test_ram_mb_positive(self, small_expert):
        ram = bench._ram_mb(small_expert)
        assert ram > 0

    def test_qa_exact_recall(self):
        pairs = [("what is x", "x is the answer")] * 4
        e = Expert.from_qa_pairs(pairs, domain="t", D=5000, seed=0)
        recall, n = bench._qa_exact_recall(e, [("what is x", "x is the answer")])
        assert n == 1
        assert 0.0 <= recall <= 1.0

    def test_query_latency_positive(self, small_expert):
        lat = bench._measure_query_latency(small_expert, ["what is"], k=3)
        assert lat > 0


class TestScalingPresets:
    def test_scales_defined(self):
        from ensemble.expert import SCALES
        assert "1b" in SCALES
        assert SCALES["1b"]["D"] == 100_000
        assert SCALES["small"]["D"] == 10_000