132 lines · 5.0 KB · Python
RawDownload
from __future__ import annotations import jsonimport osimport tempfileimport unittestfrom pathlib import Pathfrom unittest.mock import patch from zed_eval import run_index  class RunIndexTests(unittest.TestCase):    def with_index_path(self, path: Path):        return patch.dict(os.environ, {"AGENT_EVALS_RUN_INDEX": str(path)})     def test_record_lookup_and_recent(self) -> None:        with tempfile.TemporaryDirectory() as temporary_directory:            path = Path(temporary_directory) / "nested" / "run-index.json"            with self.with_index_path(path):                run_index.record_run(                    {                        "run_id": "run-1",                        "namespace": "alice",                        "experiment_name": "rf",                        "volume_name": "custom-volume",                        "agent_model": "sonnet-4.6",                        "created_at": "2026-01-01T00:00:00+00:00",                    }                )                run_index.record_run(                    {                        "run_id": "run-2",                        "namespace": "alice",                        "experiment_name": "qna",                        "volume_name": None,                    }                )                 self.assertEqual(run_index.lookup("run-1")["volume"], "custom-volume")                self.assertEqual(                    [entry["run_id"] for entry in run_index.recent(2)],                    ["run-2", "run-1"],                )                self.assertEqual(run_index.most_recent()["run_id"], "run-2")                self.assertNotIn("volume", run_index.lookup("run-2"))     def test_record_refreshes_existing_run(self) -> None:        with tempfile.TemporaryDirectory() as temporary_directory:            path = Path(temporary_directory) / "run-index.json"            with self.with_index_path(path):                run_index.record_run(                    {                        "run_id": "run-1",                        "namespace": "alice",                        "experiment_name": "rf",                    }                )                run_index.record_run(                    {                        "run_id": "run-1",                        "namespace": "bob",                        "experiment_name": "tw",                    }                )                 entries = run_index.recent(10)                self.assertEqual(len(entries), 1)                self.assertEqual(entries[0]["namespace"], "bob")                self.assertEqual(entries[0]["experiment_name"], "tw")     def test_corrupt_or_malformed_index_is_ignored(self) -> None:        with tempfile.TemporaryDirectory() as temporary_directory:            path = Path(temporary_directory) / "run-index.json"            with self.with_index_path(path):                path.write_text("not json")                self.assertIsNone(run_index.lookup("run-1"))                self.assertEqual(run_index.recent(), [])                 path.write_text(                    json.dumps(                        {                            "runs": [                                {"run_id": "missing-location"},                                ["not", "a", "dict"],                                {                                    "run_id": "run-1",                                    "namespace": "alice",                                    "experiment_name": "rf",                                },                            ]                        }                    )                )                self.assertIsNone(run_index.lookup("missing-location"))                self.assertEqual(run_index.lookup("run-1")["experiment_name"], "rf")     def test_caps_entries(self) -> None:        with tempfile.TemporaryDirectory() as temporary_directory:            path = Path(temporary_directory) / "run-index.json"            with self.with_index_path(path), patch.object(run_index, "MAX_ENTRIES", 2):                for index in range(3):                    run_index.record_run(                        {                            "run_id": f"run-{index}",                            "namespace": "alice",                            "experiment_name": "rf",                        }                    )                 self.assertEqual(                    [entry["run_id"] for entry in run_index.recent(10)],                    ["run-2", "run-1"],                )                self.assertIsNone(run_index.lookup("run-0"))     def test_zero_limit_returns_no_entries(self) -> None:        with tempfile.TemporaryDirectory() as temporary_directory:            path = Path(temporary_directory) / "run-index.json"            with self.with_index_path(path):                run_index.record_run(                    {                        "run_id": "run-1",                        "namespace": "alice",                        "experiment_name": "rf",                    }                )                self.assertEqual(run_index.recent(0), [])  if __name__ == "__main__":    unittest.main()