arena/tests/test_probe.py

from __future__ import annotations

import re
import subprocess
import sys
import tomllib
import unittest
from pathlib import Path


ROOT = Path(__file__).resolve().parents[1]
ROW = re.compile(r"^(?P<name>\S(?:.*\S)?)\s{2,}(?:-?\d+\.\d+\s*){4}\s+(?P<verdict>.+?)$")

# The per-component block, which is indented by two spaces precisely so that ROW — anchored
# at `^\S` — cannot swallow it. Four numbers there, five here.
PART = re.compile(
    r"^  (?P<env>\S+)\s+(?P<component>\S+)\s+"
    r"(?P<floor>-?\d+\.\d+)\s+(?P<below>-?\d+\.\d+)\s+"
    r"(?P<oracle>-?\d+\.\d+)\s+(?P<ceiling>-?\d+\.\d+)\s+(?P<verdict>.+?)$"
)

# Reward components that nothing below the oracle earns a fraction of. `probe.py` warns and
# names them; this is what stops the warning becoming wallpaper.
#
# Both of these were investigated by replaying `outputs/run-20260821-1401` through a
# candidate gate (`tools/regate.py`), which is the only measurement that separates a dead
# gate from a merely binary one, and both were left as they are on purpose:
#
#   redaction-pressure/gate  fired 0/32. Every margin loose enough to fire on the measured
#                            population also pays a four-rule ruleset on five seeds in six.
#                            Genuinely hard, not mis-thresholded. See `Outcome.clean`.
#   fault-localisation/gate  fired 29/32 — 0.9062, the healthiest gate in the repository.
#                            The ladder simply has no rung that gets two of three fields
#                            right and clears it, because the gate wants all three.
#
# `bot-detection/gate` and `schema-migration/gate` were on this list and are not any more:
# they came back 0/32 and were given margins. It only ever shrinks, and a name may only be
# deleted alongside the measurement that justifies it.
STEP_AT_ORACLE = frozenset({
    ("redaction-pressure", "gate"),
    ("fault-localisation", "gate"),
})

# The environments that are allowed to carry no probe yet, because their taskset layer is not
# written. Everything discovered and NOT in here MUST be gated.
#
# This list is the floor, and it is why it exists rather than a bare count: without it, both
# loops below are vacuous when nothing is gated, and CI goes green while gating zero
# environments. Commenting out a single `@probes(...)` decorator used to pass — the degated
# name simply moved into the warning line, which already legitimately carries five names.
#
# It only ever shrinks. Delete a name here in the same commit that lands its probe.
PENDING_PROBES = frozenset({
    "tera-crow-nav",
    "tera-drive-101",
    "tera-office-nav",
    "tera-california-flight",
})


def discovered_tasksets() -> set[str]:
    """The same denominator `probe.py:discover()` uses, derived independently of it.

    Deriving it a second time here rather than importing `probe` is the point: if discovery
    and the manifests ever disagree, this test is what says so. A directory with no
    `pyproject.toml` is not an environment to either of us.
    """
    names: set[str] = set()
    for manifest in sorted((ROOT / "environments").glob("*/pyproject.toml")):
        data = tomllib.loads(manifest.read_text(encoding="utf-8"))
        declared = data.get("tool", {}).get("arena", {}).get("tasksets")
        names.update(declared or [data["project"]["name"]])
    return names


class RootProbeIntegrationTests(unittest.TestCase):
    def test_public_root_probe_gates_every_environment(self) -> None:
        completed = subprocess.run(
            [sys.executable, "probe.py"],
            cwd=ROOT,
            text=True,
            capture_output=True,
            # The probe sweeps constants and runs multi-block ladders, and the interactive
            # environments will make it dearer still. A tight timeout here only ever fails a
            # slow machine, never a wrong reward.
            timeout=180,
            check=False,
        )
        detail = f"\nstdout:\n{completed.stdout}\nstderr:\n{completed.stderr}"
        self.assertEqual(completed.returncode, 0, msg=f"root probe failed:{detail}")

        discovered = discovered_tasksets()
        self.assertTrue(discovered, msg="no environment manifests found — the scan is empty")

        rows = {}
        for line in completed.stdout.splitlines()[1:]:
            match = ROW.match(line.rstrip())
            if match:
                rows[match["name"]] = match["verdict"].strip()

        for name, verdict in rows.items():
            with self.subTest(environment=name):
                self.assertIn(name, discovered, msg=f"probe gated {name}, no manifest declares it")
                self.assertEqual(verdict, "ok", msg=f"{name} is not ok:{detail}")

        # Anything discovered and not gated has to be named in the warning. That warning is
        # temporary — it becomes a hard exit 1 in probe.py once the last unregistered
        # environment lands its probe — and this assertion holds either way.
        ungated = discovered - set(rows)
        for name in sorted(ungated):
            with self.subTest(environment=name):
                self.assertIn(name, completed.stderr,
                              msg=f"{name} is neither gated nor reported as ungated:{detail}")

        # The floor. Being named in the warning is not a licence to stop being gated: an
        # environment may only be ungated if it is on the pending list. Without this, removing
        # a `@probes(...)` decorator moves the name from the table into the warning and every
        # assertion above stays vacuously true.
        self.assertLessEqual(
            ungated, PENDING_PROBES,
            msg=f"these were gated and no longer are: {sorted(ungated - PENDING_PROBES)}{detail}",
        )
        self.assertTrue(rows, msg=f"the probe gated nothing at all:{detail}")

        # --- and the same again, one reward component at a time -------------------------
        # Everything above passed for a month while four `gate` components scored exactly
        # 0.000 mean and 0.000 max over 32 real rollouts, because a blended `oracle 1.000`
        # is exactly the thing that cannot show a constant inside it.
        parts = {}
        for line in completed.stdout.splitlines():
            match = PART.match(line.rstrip())
            if match:
                parts[(match["env"], match["component"])] = match

        self.assertTrue(parts, msg=f"the probe printed no per-component block:{detail}")
        for name in rows:
            with self.subTest(environment=name):
                self.assertTrue(
                    any(env == name for env, _ in parts),
                    msg=f"{name} is gated but no component of it is reported:{detail}",
                )

        for (env, component), match in sorted(parts.items()):
            with self.subTest(environment=env, component=component):
                verdict = match["verdict"].strip()
                self.assertFalse(
                    verdict.startswith("FLAT"),
                    msg=f"{env}/{component} is constant across the ladder:{detail}",
                )
                # House rule 3, per term rather than per environment. A blended floor of
                # 0.000 is compatible with one component paying for inaction and another
                # going negative to cancel it; this is the form that is not.
                self.assertLessEqual(
                    float(match["floor"]), 1e-9,
                    msg=f"{env}/{component} pays {match['floor']} for the worst rung on the "
                        f"ladder — inaction is being paid for a component:{detail}",
                )
                stepped = verdict.startswith("step@oracle")
                if (env, component) in STEP_AT_ORACLE:
                    self.assertTrue(
                        stepped,
                        msg=f"{env}/{component} now carries a gradient below the oracle — "
                            f"delete it from STEP_AT_ORACLE:{detail}",
                    )
                else:
                    self.assertFalse(
                        stepped,
                        msg=f"{env}/{component} earns nothing below the oracle and is not "
                            f"on the STEP_AT_ORACLE list. Replay the traces through it with "
                            f"tools/regate.py before adding it:{detail}",
                    )


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