tests/test_dataset.py

"""Dataset regression tests — DATA-01..06."""
from __future__ import annotations

import hashlib
import json
import os
from collections import Counter
from pathlib import Path

import pytest


# ---------------------------------------------------------------------------
# Fixture: load generated JSONL files
# ---------------------------------------------------------------------------

@pytest.fixture(scope="session")
def generated_data():
    base = Path("data")
    train_path = base / "train.jsonl"
    eval_path = base / "eval.jsonl"
    sft_path = base / "sft_traces.jsonl"
    for p in [train_path, eval_path, sft_path]:
        assert p.exists(), f"Generated file missing: {p}. Run `python -m data.generate` first."
    train = [json.loads(l) for l in train_path.open(encoding="utf-8")]
    eval_ = [json.loads(l) for l in eval_path.open(encoding="utf-8")]
    sft = [json.loads(l) for l in sft_path.open(encoding="utf-8")]
    return {"train": train, "eval": eval_, "sft": sft}


# ---------------------------------------------------------------------------
# Test 1: counts
# ---------------------------------------------------------------------------

def test_train_count_and_eval_count(generated_data):
    assert len(generated_data["train"]) == 1000
    assert len(generated_data["eval"]) == 200


# ---------------------------------------------------------------------------
# Test 2: mix ratios
# ---------------------------------------------------------------------------

def test_mix_ratios_within_tolerance(generated_data):
    rows = generated_data["train"]
    c = Counter(r["task_type"] for r in rows)
    total = len(rows)
    expected = {"niah": 0.4, "multi_needle": 0.3, "extractive": 0.2, "counting": 0.1}
    for task_type, target in expected.items():
        actual = c[task_type] / total
        assert abs(actual - target) <= 0.02, (
            f"{task_type}: expected ~{target:.2f}, got {actual:.3f}"
        )


# ---------------------------------------------------------------------------
# Test 3: trivial floor (DATA-04)
# ---------------------------------------------------------------------------

def test_trivial_floor_satisfied(generated_data):
    rows = generated_data["train"]
    share = sum(1 for r in rows if r["difficulty"] == "trivial") / len(rows)
    assert share >= 0.05, f"Trivial share {share:.3f} < 0.05"


# ---------------------------------------------------------------------------
# Test 4: train/eval disjoint (DATA-02)
# ---------------------------------------------------------------------------

def test_train_eval_disjoint(generated_data):
    train_ids = {r["task_id"] for r in generated_data["train"]}
    eval_ids = {r["task_id"] for r in generated_data["eval"]}
    assert not (train_ids & eval_ids), "Train and eval task_ids overlap"

    train_pairs = {(r["task_type"], r["seed"]) for r in generated_data["train"]}
    eval_pairs = {(r["task_type"], r["seed"]) for r in generated_data["eval"]}
    assert not (train_pairs & eval_pairs), "Train and eval (task_type, seed) pairs overlap"


# ---------------------------------------------------------------------------
# Test 5: schema keys (DATA-05)
# ---------------------------------------------------------------------------

def test_schema_keys(generated_data):
    required = {"task_id", "task_type", "difficulty", "context_length",
                "prompt", "context", "gold_answer", "seed"}
    for split in ("train", "eval"):
        for i, row in enumerate(generated_data[split]):
            assert required.issubset(row.keys()), (
                f"{split}[{i}] missing keys: {required - row.keys()}"
            )


# ---------------------------------------------------------------------------
# Test 6: no gold_answer verbatim leak (REW-06 length/leak gate)
# ---------------------------------------------------------------------------

def test_no_gold_answer_verbatim_leak(generated_data):
    """Gold answer should appear ONLY as part of the intentional needle/fact.

    We check that the number of occurrences is exactly 1 (the needle itself).
    Multiple occurrences would indicate the filler accidentally contains the answer,
    which would enable trivial shortcut exploits (REW-06 length/leak gate).
    """
    violations = []
    for split in ("train", "eval"):
        for row in generated_data[split]:
            gold = row["gold_answer"].lower().strip()
            if not gold:
                continue
            # Count occurrences: exactly 1 is expected (the needle fact itself)
            count = row["context"].lower().count(gold)
            if count > 1:
                violations.append((row["task_id"], count))
    assert not violations, f"Gold answer found >1x in context for: {violations[:5]}"



# ---------------------------------------------------------------------------
# Test 7: SFT traces chat format (DATA-06)
# ---------------------------------------------------------------------------

def test_sft_traces_chat_format(generated_data):
    sft = generated_data["sft"]
    assert len(sft) >= 450, f"Only {len(sft)} SFT traces (need >= 450)"
    valid_roles = {"system", "user", "assistant", "tool"}
    for i, trace in enumerate(sft):
        assert "messages" in trace, f"trace[{i}] missing 'messages'"
        messages = trace["messages"]
        assert len(messages) >= 3, f"trace[{i}] has only {len(messages)} messages"
        roles = {m["role"] for m in messages}
        assert roles.issubset(valid_roles), f"trace[{i}] invalid roles: {roles}"
        # Every trace must have at least one assistant message with <answer>
        has_answer = any(
            m["role"] == "assistant" and "<answer>" in m.get("content", "")
            for m in messages
        )
        assert has_answer, f"trace[{i}] has no assistant <answer> message"


# ---------------------------------------------------------------------------
# Test 8: byte-determinism (marked slow) (DATA-03)
# ---------------------------------------------------------------------------

@pytest.mark.slow
def test_generation_is_byte_deterministic(tmp_path, monkeypatch):
    """Re-running generate_all with same seeds produces byte-identical output."""
    monkeypatch.delenv("ANTHROPIC_API_KEY", raising=False)

    from data.generate import generate_all

    generate_all(out_dir=tmp_path / "run1", seeds_path="data/seeds.json")
    generate_all(out_dir=tmp_path / "run2", seeds_path="data/seeds.json")

    for fname in ("train.jsonl", "eval.jsonl", "sft_traces.jsonl"):
        h1 = hashlib.sha256((tmp_path / "run1" / fname).read_bytes()).hexdigest()
        h2 = hashlib.sha256((tmp_path / "run2" / fname).read_bytes()).hexdigest()
        assert h1 == h2, f"{fname}: non-deterministic output (run1 != run2)"