bench/tests/test_data.py

"""Tests over the eval data itself.

Data rots differently from code: nothing crashes when a sample is quietly malformed, an id is
reused, or a `test` sample creeps into a training split. The score just moves, and the move
looks exactly like a real result. These are the checks that fail loudly instead.

    python3 -m unittest discover -s tests -v
"""

from __future__ import annotations

import re
import sys
import unittest
from collections import Counter
from pathlib import Path

ROOT = Path(__file__).resolve().parent.parent
sys.path.insert(0, str(ROOT))

from kbench.schema import read_records  # noqa: E402

DATA = ROOT / "data"
FAMILIES = sorted(DATA.glob("*/*.jsonl"))

# Anything that must never reach a third-party judge or baseline model. The authoring rules
# say repo privacy is defence in depth, not the control — this is the control.
FORBIDDEN = [
    (r"\b115\b", "customer identifier"),
    (r"workstation|tailnet|cloud-[123]-|your-second-node|your-node", "internal host"),
    (r"\b(?:\d{1,3}\.){3}\d{1,3}\b", "IP address"),
    (r"sk-[A-Za-z0-9]{16,}|ghp_[A-Za-z0-9]{16,}", "API key"),
    (r"-----BEGIN [A-Z ]*PRIVATE KEY", "private key"),
]


class DataIntegrity(unittest.TestCase):
    def test_every_family_parses_and_validates(self):
        self.assertTrue(FAMILIES, "no data files found at all")
        for path in FAMILIES:
            with self.subTest(family=path.name):
                self.assertTrue(read_records(path), f"{path.name} is empty")

    def test_ids_are_unique_within_a_family(self):
        # Per-sample regression tracking keys on the id. A duplicate silently makes one of the
        # two invisible in every future comparison.
        for path in FAMILIES:
            ids = [r["id"] for _, r in read_records(path)]
            dupes = [i for i, n in Counter(ids).items() if n > 1]
            self.assertEqual(dupes, [], f"{path.name} reuses ids: {dupes}")

    def test_no_two_samples_share_a_prompt(self):
        # Unique ids are not enough. Re-running a seeder appends a second copy of every
        # sample under fresh ids, which passes the id check, doubles the family, and silently
        # double-weights whatever those samples measure. This has already happened once.
        for path in FAMILIES:
            inputs = [r["input"] for _, r in read_records(path)]
            dupes = [i for i, n in Counter(inputs).items() if n > 1]
            with self.subTest(family=path.name):
                self.assertEqual(
                    [d[:60] for d in dupes],
                    [],
                    f"{path.name} contains the same prompt more than once",
                )

    def test_ids_match_their_family_and_file(self):
        for path in FAMILIES:
            for lineno, r in read_records(path):
                with self.subTest(sample=r["id"]):
                    self.assertTrue(
                        r["id"].startswith(f"{r.get('family')}/"),
                        f"{path.name}:{lineno} id {r['id']} does not match family {r.get('family')}",
                    )

    def test_regex_targets_actually_compile(self):
        # An invalid regex target does not raise at author time; it fails at grading time,
        # mid-run, after the model has already been paid for.
        for path in FAMILIES:
            for lineno, r in read_records(path):
                if r["scorer"] != "regex":
                    continue
                with self.subTest(sample=r["id"]):
                    try:
                        re.compile(r["target"])
                    except re.error as exc:
                        self.fail(f"{path.name}:{lineno} {r['id']} bad regex: {exc}")

    def test_no_sensitive_content_reaches_a_third_party_model(self):
        for path in FAMILIES:
            for lineno, r in read_records(path):
                haystack = " ".join(
                    str(r.get(k) or "") for k in ("input", "target", "rubric")
                )
                for pattern, label in FORBIDDEN:
                    with self.subTest(sample=r["id"], rule=label):
                        self.assertIsNone(
                            re.search(pattern, haystack, re.I),
                            f"{path.name}:{lineno} {r['id']} contains a {label}",
                        )

    def test_rubric_samples_are_rare(self):
        # Each rubric sample adds judge cost per run and judge drift between runs, which
        # confounds exactly the checkpoint comparisons this bench exists to make.
        for path in FAMILIES:
            records = [r for _, r in read_records(path)]
            rubrics = [r for r in records if r["scorer"] == "rubric"]
            with self.subTest(family=path.name):
                self.assertLessEqual(
                    len(rubrics),
                    max(1, len(records) // 4),
                    f"{path.name}: {len(rubrics)}/{len(records)} samples are judge-graded",
                )


class ContaminationProbes(unittest.TestCase):
    """The canary tier is the one place where passing is bad."""

    PATH = DATA / "canary" / "contamination.jsonl"

    def setUp(self):
        if not self.PATH.exists():
            self.skipTest("no canary family")
        self.records = [r for _, r in read_records(self.PATH)]

    def _carriers(self):
        """Probes that put the GUID into the traffic. Their answer is in their own prompt."""
        return [r for r in self.records if r["target"] in r["input"]]

    def _detectors(self):
        """Probes that ask for the GUID cold. Only a trained-on model can answer."""
        return [r for r in self.records if r["target"] not in r["input"]]

    def test_all_probes_target_one_shared_canary(self):
        # Carrier and detectors must chase the SAME string. A detector for a GUID no carrier
        # ever transmitted is unanswerable by construction — it can never fire, so it proves
        # nothing while looking like a control. That was the first version of this file.
        targets = {r["target"] for r in self.records}
        self.assertEqual(len(targets), 1, f"probes chase different canaries: {targets}")

    def test_exactly_one_carrier_transmits_the_guid(self):
        carriers = self._carriers()
        self.assertEqual(
            len(carriers), 1, "need exactly one carrier: the probe that leaks the GUID on purpose"
        )
        self.assertIn("LBCANARY-", carriers[0]["input"])

    def test_at_least_two_detectors_withhold_the_guid(self):
        detectors = self._detectors()
        self.assertGreaterEqual(len(detectors), 2, "need at least two cold detectors")
        for r in detectors:
            with self.subTest(sample=r["id"]):
                # A detector may name the canary, and must not contain the answer.
                self.assertNotIn(r["target"], r["input"])

    def test_a_detector_never_leaks_the_suffix_it_is_testing_for(self):
        # A prefix hint is fine and useful; the tail must be absent or the probe grades itself.
        for r in self._detectors():
            suffix = r["target"].split("-", 2)[-1]
            with self.subTest(sample=r["id"]):
                self.assertNotIn(suffix, r["input"], "detector prompt contains the answer suffix")

    def test_role_tags_agree_with_the_structural_role(self):
        # Two sources of truth for the same fact: the `tags` field, which the task code reads
        # to decide what to score, and whether the prompt contains its own answer, which is
        # what actually makes a probe a carrier. If they drift, the task scores the wrong set
        # and the disagreement is invisible at runtime.
        for r in self.records:
            tags = (r.get("metadata") or {}).get("tags", [])
            structural = "carrier" if r["target"] in r["input"] else "detector"
            with self.subTest(sample=r["id"]):
                self.assertIn(
                    structural,
                    tags,
                    f"{r['id']} is structurally a {structural} but is tagged {tags}",
                )


class CanaryTaskScoresOnlyDetectors(unittest.TestCase):
    """The carrier must be loaded and then not scored.

    Regression: the task scored every probe including the carrier. A carrier states the GUID
    in its own prompt and asks for it back, so it always passes — which pinned the canary at
    >= 1/n for every model alive, reported CONTAMINATED, and nulled signal_score. The data
    said `carrier` and carried a note reading "it always passes"; the task never looked.
    """

    def setUp(self):
        # The published tree ships neither family: the signal set would be contaminated by
        # publication, and a published canary GUID would enter public corpora and fire
        # against every model that ever read the repo. So both absences are expected here.
        if not (DATA / "canary" / "contamination.jsonl").exists():
            self.skipTest("no canary family (expected in the published tree)")
        try:
            from kbench.tasks.canary import canary_task
        except ImportError as exc:  # inspect-ai is not needed for the data-only tests
            self.skipTest(f"inspect-ai unavailable: {exc}")
        self.canary_task = canary_task

    def test_carrier_is_excluded_from_the_scored_dataset(self):
        task = self.canary_task("contamination")
        scored = list(task.dataset)
        self.assertTrue(scored, "canary task scored nothing")
        for sample in scored:
            with self.subTest(sample=sample.id):
                self.assertNotIn(
                    str(sample.target),
                    str(sample.input),
                    f"{sample.id} contains its own answer, so it always passes; "
                    "carriers must not be scored",
                )

    def test_every_scored_probe_is_tagged_detector(self):
        for sample in self.canary_task("contamination").dataset:
            with self.subTest(sample=sample.id):
                self.assertIn("detector", (sample.metadata or {}).get("tags", []))


if __name__ == "__main__":
    unittest.main()