tenant.openagents/omega

1from __future__ import annotations
2
3import argparse
4import contextlib
5import io
6import unittest
7from unittest.mock import patch
8
9from zed_eval import config
10from zed_eval.cli import build_parser
11from zed_eval.launch import prepare_benchmark_suite, prepare_shared_build, prepare_suite
12
13
14class LaunchPlanTests(unittest.TestCase):
15    def parse(self, argv: list[str]) -> argparse.Namespace:
16        return build_parser().parse_args(argv)
17
18    def test_swe_atlas_defaults_judges_by_part(self) -> None:
19        args = self.parse(
20            [
21                "swe-atlas",
22                "--parts",
23                "qna,rf",
24                "--build",
25                "bld-test",
26                "--plan",
27                "--allow-untracked",
28                "--n-tasks",
29                "1",
30                "--yes",
31            ]
32        )
33
34        prepared = prepare_suite(args)
35
36        self.assertEqual([part for part, *_ in prepared], ["qna", "rf"])
37        self.assertEqual(prepared[0][1]["build_id"], "bld-test")
38        self.assertEqual(prepared[0][2]["build_id"], "bld-test")
39        self.assertEqual(prepared[1][1]["build_id"], "bld-test")
40        self.assertIsNone(prepared[1][2])
41        self.assertEqual(prepared[0][1]["judge_preset"], "deepseek-v4-pro")
42        self.assertEqual(prepared[1][1]["judge_preset"], "kimi-k2.7-code")
43        self.assertEqual(prepared[0][1]["agent_model"], config.DEFAULT_MODEL)
44
45    def test_baseten_model_preset_generates_provider_json(self) -> None:
46        args = self.parse(
47            [
48                "swe-atlas",
49                "--parts",
50                "tw",
51                "--build",
52                "bld-test",
53                "--plan",
54                "--allow-untracked",
55                "--model",
56                "baseten:kimi-k2.7-code",
57                "--yes",
58            ]
59        )
60
61        prepared = prepare_suite(args)
62        run_request = prepared[0][1]
63
64        self.assertEqual([part for part, *_ in prepared], ["tw"])
65        self.assertEqual(run_request["benchmark"]["id"], "swe-atlas-tw")
66        self.assertEqual(
67            run_request["agent_model"], "baseten/moonshotai/Kimi-K2.7-Code"
68        )
69        self.assertIn(
70            "moonshotai/Kimi-K2.7-Code", run_request["openai_compatible_provider_json"]
71        )
72        self.assertIn(
73            "https://inference.baseten.co/v1",
74            run_request["openai_compatible_provider_json"],
75        )
76
77    def test_interactive_suite_can_include_non_swe_atlas_benchmarks(self) -> None:
78        args = self.parse(
79            [
80                "swe-atlas",
81                "--build",
82                "bld-test",
83                "--plan",
84                "--allow-untracked",
85                "--yes",
86            ]
87        )
88
89        prepared = prepare_benchmark_suite(
90            args, ["qna", "terminal-bench-2.1", "deepswe"]
91        )
92
93        self.assertEqual(
94            [
95                run_request["benchmark"]["id"]
96                for _label, run_request, _build in prepared
97            ],
98            ["swe-atlas-qna", "terminal-bench-2.1", "deepswe"],
99        )
100        self.assertEqual(
101            [label for label, *_ in prepared], ["qna", "terminal-bench-2.1", "deepswe"]
102        )
103        self.assertEqual(prepared[0][1]["judge_preset"], "deepseek-v4-pro")
104        self.assertIsNone(prepared[1][1]["judge_preset"])
105        self.assertIsNone(prepared[2][1]["judge_preset"])
106        self.assertEqual(prepared[0][1]["build_id"], "bld-test")
107        self.assertEqual(prepared[1][1]["build_id"], "bld-test")
108        self.assertEqual(prepared[2][1]["build_id"], "bld-test")
109        self.assertEqual(prepared[0][2]["build_id"], "bld-test")
110        self.assertIsNone(prepared[1][2])
111        self.assertIsNone(prepared[2][2])
112
113    def test_explicit_existing_build_skips_source_preparation(self) -> None:
114        args = self.parse(["run", "rf", "--build", "custom-build"])
115
116        with (
117            contextlib.redirect_stdout(io.StringIO()),
118            patch("zed_eval.launch.build_ready_on_volume", return_value=True),
119            patch("zed_eval.launch.resolve_source") as resolve_source,
120        ):
121            build_id, build_request = prepare_shared_build(args)
122
123        self.assertEqual(build_id, "custom-build")
124        self.assertIsNone(build_request)
125        resolve_source.assert_not_called()
126
127    def test_explicit_missing_build_uses_from_source(self) -> None:
128        args = self.parse(
129            ["run", "rf", "--build", "custom-build", "--from", "v0.210.0"]
130        )
131
132        with (
133            patch("zed_eval.launch.build_ready_on_volume", return_value=False),
134            patch("zed_eval.source.resolve_remote_ref", return_value="a" * 40),
135        ):
136            build_id, build_request = prepare_shared_build(args)
137
138        self.assertEqual(build_id, "custom-build")
139        self.assertEqual(build_request["build_id"], "custom-build")
140        self.assertEqual(build_request["source"]["base_sha"], "a" * 40)
141        self.assertEqual(build_request["source"]["base_ref"], "v0.210.0")
142        self.assertTrue(build_request["source"]["clean_source"])
143
144
145if __name__ == "__main__":
146    unittest.main()
147