tests/test_eval_harness.py

from __future__ import annotations

import argparse
import asyncio
import csv
import json
import tempfile
import unittest
from pathlib import Path
from types import SimpleNamespace
from unittest import mock

from evals.check_triggers import check_run
from evals.report import write_cost_curves
from evals.run_compaction_experiment import (
    MIGRATION_FILENAME,
    TRIGGER_VALIDATION_BASELINE_SOURCE_SHA256,
    _bundle_progress,
    _import_compatible_completed_pairs,
    _is_retryable_turn_error,
    _prune_to_common_completed,
    _run_turn_with_retries,
    _run_session_trial,
    _schedule_pending,
    _write_run_status,
)
from evals.run_battery import (
    experiment_cache_user_id,
    validate_experiment_result,
    write_or_validate_run_config,
)


def experiment_args(battery: Path, out: Path) -> argparse.Namespace:
    return argparse.Namespace(
        battery=str(battery),
        preset="exp_fh_raw",
        model="deepseek:deepseek-v4-flash",
        trials=1,
        arm_concurrency=1,
        pair_concurrency=1,
        max_pairs_this_invocation=0,
        first_pair_id="",
        import_completed_from="",
        out=str(out),
        limit=0,
        ids=[],
        tags=[],
        concurrency=1,
        scope_sources=False,
        enable_tools=[],
        disable_kb=False,
        no_tools=False,
        retrieval_budget=0,
        retriever="",
        langsmith=False,
    )


class ImmutableManifestTests(unittest.TestCase):
    def test_status_file_is_atomic_machine_readable_and_includes_fatal_error(
        self,
    ) -> None:
        with tempfile.TemporaryDirectory() as directory:
            root = Path(directory)
            args = experiment_args(root / "battery.jsonl", root)
            args.presets = ["exp_fh_raw"]
            error = RuntimeError("guard exceeded")
            path = _write_run_status(
                root,
                state="failed",
                args=args,
                progress={"completed_pairs_all_arms": 4},
                error=error,
                formatted_traceback="traceback text",
            )
            status = json.loads(path.read_text())
            self.assertEqual(status["state"], "failed")
            self.assertEqual(status["fatal_error"]["type"], "RuntimeError")
            self.assertEqual(status["fatal_error"]["message"], "guard exceeded")
            self.assertEqual(status["progress"]["completed_pairs_all_arms"], 4)
            self.assertFalse(path.with_suffix(".json.tmp").exists())

    def test_bundle_progress_reports_only_pairs_complete_in_every_arm(self) -> None:
        class Sink:
            def __init__(self, path: Path) -> None:
                self.path = path

        with tempfile.TemporaryDirectory() as directory:
            root = Path(directory)
            paths = {name: root / f"{name}.jsonl" for name in ("a", "b")}
            complete = [
                {"unit_id": "s1", "trial": 1, "error": None},
                {"unit_id": "s1", "trial": 1, "error": None},
            ]
            paths["a"].write_text("".join(json.dumps(row) + "\n" for row in complete))
            paths["b"].write_text(
                "".join(json.dumps(row) + "\n" for row in complete[:1])
            )
            progress = _bundle_progress(
                {name: Sink(path) for name, path in paths.items()}, {"s1": 2}
            )
            self.assertEqual(progress["arms"]["a"]["completed_pairs"], 1)
            self.assertEqual(progress["arms"]["b"]["completed_pairs"], 0)
            self.assertEqual(progress["completed_pairs_all_arms"], 0)

    def test_resume_accepts_exact_match_and_rejects_drift(self) -> None:
        with tempfile.TemporaryDirectory() as directory:
            root = Path(directory)
            battery = root / "battery.jsonl"
            battery.write_text('{"session_id":"s","turns":["q"]}\n')
            out = root / "run"
            out.mkdir()
            args = experiment_args(battery, out)
            first = write_or_validate_run_config(out, args)
            (out / "bundles.jsonl").write_text("{}\n")
            self.assertEqual(write_or_validate_run_config(out, args), first)

            args.retrieval_budget = 30_000
            with self.assertRaises(SystemExit):
                write_or_validate_run_config(out, args)

            config = json.loads((out / "run_config.json").read_text())
            self.assertEqual(config["_fingerprint"], first)
            self.assertIn("source_tree_sha256", config["_manifest"])
            self.assertIn("pricing_snapshot_usd_per_million", config["_manifest"])

    def test_operational_staging_controls_do_not_change_fingerprint(self) -> None:
        with tempfile.TemporaryDirectory() as directory:
            root = Path(directory)
            battery = root / "battery.jsonl"
            battery.write_text('{"session_id":"s","turns":["q"]}\n')
            out = root / "run"
            out.mkdir()
            args = experiment_args(battery, out)
            first = write_or_validate_run_config(out, args)
            (out / "bundles.jsonl").write_text("{}\n")
            args.max_pairs_this_invocation = 1
            args.first_pair_id = "s"
            args.import_completed_from = str(root / "old-run")
            self.assertEqual(write_or_validate_run_config(out, args), first)

    def test_cache_ids_are_stable_and_arm_isolated(self) -> None:
        first = experiment_cache_user_id("exp_fh_raw", "session", 1)
        self.assertEqual(first, experiment_cache_user_id("exp_fh_raw", "session", 1))
        self.assertNotEqual(
            first, experiment_cache_user_id("exp_c200_raw", "session", 1)
        )
        self.assertNotEqual(
            first,
            experiment_cache_user_id("exp_fh_raw", "session", 1, namespace="new-run"),
        )
        self.assertRegex(first, r"^[a-zA-Z0-9_-]+$")

    def test_lockstep_resume_keeps_only_units_complete_in_every_arm(self) -> None:
        with tempfile.TemporaryDirectory() as directory:
            root = Path(directory)
            arm_a = root / "a.jsonl"
            arm_b = root / "b.jsonl"
            complete = [
                {"unit_id": "s1", "trial": 1, "turn_index": 0, "error": None},
                {"unit_id": "s1", "trial": 1, "turn_index": 1, "error": None},
            ]
            partial = [{"unit_id": "s2", "trial": 1, "turn_index": 0, "error": None}]
            arm_a.write_text(
                "".join(json.dumps(row) + "\n" for row in [*complete, *partial])
            )
            arm_b.write_text("".join(json.dumps(row) + "\n" for row in complete))
            common = _prune_to_common_completed([arm_a, arm_b], {"s1": 2, "s2": 2})
            self.assertEqual(common, {("s1", 1)})
            self.assertEqual(len(arm_a.read_text().splitlines()), 2)

    def _write_migration_fixture(
        self, root: Path, *, drift: bool = False
    ) -> tuple[Path, Path, list[str], dict[str, str]]:
        source = root / "source"
        target = root / "target"
        presets = ["exp_fh_raw", "exp_c200_raw"]
        target_fingerprints: dict[str, str] = {}
        for preset in presets:
            source_arm = source / preset
            target_arm = target / preset
            source_arm.mkdir(parents=True)
            target_arm.mkdir(parents=True)
            source_fingerprint = f"source-{preset}"
            target_fingerprint = f"target-{preset}"
            target_fingerprints[preset] = target_fingerprint
            scientific_value = "changed" if drift and preset == presets[-1] else "same"
            source_manifest = {
                "source_tree_sha256": TRIGGER_VALIDATION_BASELINE_SOURCE_SHA256,
                "git_status": "old dirty state",
                "scientific_configuration": "same",
            }
            target_manifest = {
                "source_tree_sha256": "new-source-tree",
                "git_status": "new dirty state",
                "scientific_configuration": scientific_value,
            }
            (source_arm / "run_config.json").write_text(
                json.dumps(
                    {
                        "_fingerprint": source_fingerprint,
                        "_manifest": source_manifest,
                    }
                )
            )
            (target_arm / "run_config.json").write_text(
                json.dumps(
                    {
                        "_fingerprint": target_fingerprint,
                        "_manifest": target_manifest,
                    }
                )
            )
            complete = [
                {
                    "unit_id": "s1",
                    "trial": 1,
                    "turn_index": turn,
                    "error": None,
                    "run_fingerprint": source_fingerprint,
                }
                for turn in range(2)
            ]
            partial = [
                {
                    "unit_id": "s2",
                    "trial": 1,
                    "turn_index": 0,
                    "error": None,
                    "run_fingerprint": source_fingerprint,
                }
            ]
            rows = [*complete, *partial] if preset == presets[0] else complete
            (source_arm / "bundles.jsonl").write_text(
                "".join(json.dumps(row) + "\n" for row in rows)
            )
        return source, target, presets, target_fingerprints

    def test_migration_imports_only_common_complete_pairs_with_provenance(self) -> None:
        with tempfile.TemporaryDirectory() as directory:
            root = Path(directory)
            source, target, presets, target_fingerprints = (
                self._write_migration_fixture(root)
            )
            imported = _import_compatible_completed_pairs(
                source_root=source,
                target_root=target,
                presets=presets,
                target_fingerprints=target_fingerprints,
                expected_turns={"s1": 2, "s2": 2},
                eligible_pairs={("s1", 1), ("s2", 1)},
            )
            self.assertEqual(imported, {("s1", 1)})
            for preset in presets:
                rows = [
                    json.loads(line)
                    for line in (target / preset / "bundles.jsonl")
                    .read_text()
                    .splitlines()
                ]
                self.assertEqual(len(rows), 2)
                self.assertTrue(
                    all(
                        row["run_fingerprint"] == target_fingerprints[preset]
                        for row in rows
                    )
                )
                self.assertTrue(
                    all(
                        row["migration"]["source_run_fingerprint"] == f"source-{preset}"
                        for row in rows
                    )
                )
            record = json.loads((target / MIGRATION_FILENAME).read_text())
            self.assertEqual(record["imported_pairs"], [{"trial": 1, "unit_id": "s1"}])
            self.assertEqual(
                _import_compatible_completed_pairs(
                    source_root=source,
                    target_root=target,
                    presets=presets,
                    target_fingerprints=target_fingerprints,
                    expected_turns={"s1": 2, "s2": 2},
                    eligible_pairs={("s1", 1), ("s2", 1)},
                ),
                {("s1", 1)},
            )

    def test_migration_rejects_scientific_manifest_drift(self) -> None:
        with tempfile.TemporaryDirectory() as directory:
            source, target, presets, fingerprints = self._write_migration_fixture(
                Path(directory), drift=True
            )
            with self.assertRaisesRegex(SystemExit, "scientific manifest drift"):
                _import_compatible_completed_pairs(
                    source_root=source,
                    target_root=target,
                    presets=presets,
                    target_fingerprints=fingerprints,
                    expected_turns={"s1": 2, "s2": 2},
                    eligible_pairs={("s1", 1), ("s2", 1)},
                )


class PairedConcurrencyTests(unittest.IsolatedAsyncioTestCase):
    async def test_transient_turn_failure_retries_and_records_missing_usage(
        self,
    ) -> None:
        failed = {
            "started_at": "first",
            "duration_ms": 10,
            "error": "ReadError: stream reset",
            "context_stats": None,
        }
        succeeded = {
            "started_at": "second",
            "duration_ms": 20,
            "error": None,
            "context_stats": {},
        }
        with (
            mock.patch(
                "evals.run_compaction_experiment.run_turn",
                side_effect=[failed, succeeded],
            ) as run,
            mock.patch("evals.run_compaction_experiment.asyncio.sleep") as sleep,
        ):
            result = await _run_turn_with_retries(SimpleNamespace())
        self.assertEqual(run.await_count, 2)
        sleep.assert_awaited_once_with(1.0)
        self.assertEqual(result["started_at"], "first")
        self.assertEqual(result["turn_retry_attempts"], 1)
        self.assertTrue(result["failed_attempt_usage_unavailable"])
        self.assertEqual(result["context_stats"]["turn_retry_attempts"], 1)

    async def test_permanent_turn_failure_is_not_retried(self) -> None:
        failed = {
            "started_at": "first",
            "duration_ms": 10,
            "error": "RuntimeError: invalid state",
            "context_stats": None,
        }
        with mock.patch(
            "evals.run_compaction_experiment.run_turn", return_value=failed
        ) as run:
            result = await _run_turn_with_retries(SimpleNamespace())
        self.assertEqual(run.await_count, 1)
        self.assertEqual(result["turn_retry_attempts"], 0)
        self.assertFalse(_is_retryable_turn_error(result["error"]))

    def test_staging_prioritizes_trial_one_and_limits_invocation(self) -> None:
        sessions = [
            {"session_id": "short"},
            {"session_id": "representative"},
        ]
        pending = [(session, trial) for session in sessions for trial in (1, 2, 3)]
        scheduled = _schedule_pending(
            pending,
            seed=7,
            first_pair_id="representative",
            max_pairs=1,
        )
        self.assertEqual(scheduled[0][0]["session_id"], "representative")
        self.assertEqual(scheduled[0][1], 1)

    async def test_arms_overlap_but_turn_barrier_is_preserved(self) -> None:
        class Sink:
            def __init__(self) -> None:
                self.rows: list[dict] = []

            async def write(self, rows: list[dict]) -> None:
                self.rows.extend(rows)

        active = 0
        max_active = 0
        completed_turn_zero = 0

        def build_request(args, *, query, **_kwargs):
            return SimpleNamespace(preset=args.preset, query=query)

        async def run_turn(request):
            nonlocal active, max_active, completed_turn_zero
            if request.query == "q1":
                self.assertEqual(completed_turn_zero, 2)
            active += 1
            max_active = max(max_active, active)
            await asyncio.sleep(0.01)
            active -= 1
            if request.query == "q0":
                completed_turn_zero += 1
            return {
                "thread_id": f"thread-{request.preset}",
                "answer": f"answer-{request.preset}",
                "error": None,
            }

        arm_args = {
            preset: SimpleNamespace(preset=preset)
            for preset in ("exp_fh_raw", "exp_c200_raw")
        }
        sinks = {preset: Sink() for preset in arm_args}
        with (
            mock.patch(
                "evals.run_compaction_experiment.build_request",
                side_effect=build_request,
            ),
            mock.patch(
                "evals.run_compaction_experiment.run_turn", side_effect=run_turn
            ),
            mock.patch("evals.run_compaction_experiment.validate_experiment_result"),
            mock.patch(
                "evals.run_compaction_experiment.make_bundle",
                side_effect=lambda **kwargs: {"turn_index": kwargs["turn_index"]},
            ),
        ):
            await _run_session_trial(
                session={"session_id": "s", "turns": ["q0", "q1"]},
                trial=1,
                arm_args=arm_args,
                sinks=sinks,
                seed=7,
                arm_concurrency=2,
            )

        self.assertEqual(max_active, 2)
        self.assertEqual(completed_turn_zero, 2)
        self.assertTrue(all(len(sink.rows) == 2 for sink in sinks.values()))


class ExperimentValidationTests(unittest.TestCase):
    def _result(self) -> dict:
        return {
            "error": None,
            "context_stats": {
                "llm_calls": 1,
                "cost_breakdown": {"total_usd": 0.1},
                "model_calls": [
                    {
                        "sequence": 1,
                        "model": "deepseek-v4-flash",
                        "usage_reported": True,
                        "cache_details_reported": True,
                    }
                ],
                "compaction_events": [],
            },
        }

    def test_valid_deepseek_telemetry_passes(self) -> None:
        args = argparse.Namespace(
            preset="exp_fh_raw", model="deepseek:deepseek-v4-flash"
        )
        result = self._result()
        validate_experiment_result(args, result)
        self.assertIsNone(result["error"])

    def test_provider_drift_and_old_4k_summary_fail_loudly(self) -> None:
        args = argparse.Namespace(
            preset="exp_c200_raw", model="deepseek:deepseek-v4-flash"
        )
        result = self._result()
        result["context_stats"]["model_calls"][0]["model"] = "gemini-3.5-flash"
        result["context_stats"]["compaction_events"] = [
            {
                "summary_input_untrimmed": True,
                "summary_input_tokens_approx": 4_000,
            }
        ]
        validate_experiment_result(args, result)
        self.assertIn("expected only deepseek-v4-flash", result["error"])
        self.assertIn("historical 4k cap", result["error"])

    def test_provider_reported_trigger_evidence_passes_validation(self) -> None:
        args = argparse.Namespace(
            preset="exp_c200_raw", model="deepseek:deepseek-v4-flash"
        )
        result = self._result()
        result["context_stats"]["compaction_events"] = [
            {
                "configured_trigger_tokens": 200_000,
                "pre_compaction_tokens_approx": 199_567,
                "trigger_reported_tokens": 207_336,
                "trigger_source": "provider_reported",
                "summary_input_untrimmed": True,
                "summary_input_tokens_approx": 158_014,
            }
        ]
        validate_experiment_result(args, result)
        self.assertIsNone(result["error"])


class TriggerGateTests(unittest.TestCase):
    def test_expected_trigger_checks_configuration_without_rejecting_overshoot(
        self,
    ) -> None:
        with tempfile.TemporaryDirectory() as directory:
            root = Path(directory)
            battery = root / "battery.jsonl"
            battery.write_text('{"session_id":"s","turns":["q"]}\n')
            run = root / "run"
            run.mkdir()
            bundle = {
                "unit_id": "s",
                "trial": 1,
                "turn_index": 0,
                "preset": "exp_c200_raw",
                "battery_path": str(battery),
                "error": None,
                "context_stats": {
                    "input_tokens": 1,
                    "est_cost_usd": 0,
                    "summary_messages": 1,
                    "compactions_this_turn": 1,
                    "compaction_events": [
                        {
                            "configured_trigger_tokens": 200_000,
                            "pre_compaction_tokens_approx": 199_567,
                            "trigger_reported_tokens": 207_336,
                            "summary_input_tokens_approx": 205_000,
                        }
                    ],
                },
            }
            (run / "bundles.jsonl").write_text(json.dumps(bundle) + "\n")
            self.assertTrue(
                check_run(
                    run,
                    False,
                    min_compactions=1,
                    min_summary_input=4_001,
                    expected_trigger_tokens=200_000,
                    first_pre_tokens_min=200_000,
                )
            )
            self.assertFalse(check_run(run, False, expected_trigger_tokens=300_000))


class CostCurveTests(unittest.TestCase):
    def test_csv_contains_cached_uncached_output_and_total_cost(self) -> None:
        bundles = []
        for turn, total in enumerate((0.1, 0.2)):
            bundles.append(
                {
                    "unit_id": "session",
                    "trial": 1,
                    "turn_index": turn,
                    "context_stats": {
                        "cost_breakdown": {
                            "cache_read_input_usd": total / 10,
                            "cache_miss_input_usd": total / 2,
                            "cache_creation_input_usd": 0,
                            "output_usd": total * 0.4,
                            "total_usd": total,
                        },
                        "summarization_cost_usd": 0,
                    },
                }
            )
        run = {
            "battery_type": "sessions",
            "label": "exp_fh_raw",
            "bundles": bundles,
        }
        with tempfile.TemporaryDirectory() as directory:
            out = Path(directory)
            with mock.patch("evals.report._plot_cost_curves", side_effect=ImportError):
                write_cost_curves([run], out)
            with (out / "trajectory_cost_by_turn.csv").open() as stream:
                rows = list(csv.DictReader(stream))
        self.assertEqual(len(rows), 2)
        self.assertAlmostEqual(float(rows[-1]["cumulative_total_usd"]), 0.3)
        self.assertIn("turn_cache_read_input_usd", rows[-1])
        self.assertIn("turn_cache_miss_input_usd", rows[-1])
        self.assertIn("turn_output_usd", rows[-1])


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


class DeprecatedBatteryGuardTest(unittest.TestCase):
    def test_deprecated_battery_refuses_new_runs(self) -> None:
        from evals.common import ensure_battery_not_deprecated

        with self.assertRaises(SystemExit) as ctx:
            ensure_battery_not_deprecated(
                "data/eval/battery_sessions_v2.jsonl", override=False
            )
        self.assertIn("battery_sessions_v2_1.jsonl", str(ctx.exception))

    def test_override_and_healthy_batteries_pass(self) -> None:
        from evals.common import ensure_battery_not_deprecated

        ensure_battery_not_deprecated(
            "data/eval/battery_sessions_v2.jsonl", override=True
        )
        ensure_battery_not_deprecated(
            "data/eval/battery_sessions_v1.jsonl", override=False
        )