tenant.openagents/omega

1from __future__ import annotations
2
3import argparse
4import io
5import os
6import shlex
7import subprocess
8import tarfile
9import tempfile
10import unittest
11from pathlib import Path
12from unittest.mock import patch
13
14from zed_eval import common
15from zed_eval.agent_common import (
16    add_anthropic_available_models_env,
17    eval_cli_with_log_command,
18)
19
20
21class EvalCliWithLogCommandTests(unittest.TestCase):
22    def run_logged_command(self, script: str, *, timeout_message: str | None = None):
23        with tempfile.TemporaryDirectory() as temporary_directory:
24            log_path = Path(temporary_directory) / "eval-cli.txt"
25            command = eval_cli_with_log_command(
26                ["sh", "-c", shlex.quote(script)],
27                str(log_path),
28                timeout_message=timeout_message,
29            )
30            completed = subprocess.run(
31                command,
32                shell=True,
33                capture_output=True,
34                text=True,
35            )
36            log = log_path.read_text() if log_path.exists() else ""
37            return completed, log
38
39    def test_preserves_non_timeout_exit_status_through_tee(self) -> None:
40        completed, log = self.run_logged_command("echo before-failure; exit 7")
41
42        self.assertEqual(completed.returncode, 7)
43        self.assertIn("before-failure", completed.stdout)
44        self.assertIn("before-failure", log)
45
46    def test_maps_eval_cli_timeout_exit_to_success(self) -> None:
47        completed, log = self.run_logged_command(
48            "echo partial-output; exit 2",
49            timeout_message="timeout converted",
50        )
51
52        self.assertEqual(completed.returncode, 0)
53        self.assertIn("partial-output", completed.stdout)
54        self.assertIn("timeout converted", completed.stdout)
55        self.assertIn("partial-output", log)
56        self.assertIn("timeout converted", log)
57
58
59class EnvForwardingTests(unittest.TestCase):
60    def test_anthropic_available_models_env_is_added(self) -> None:
61        env: dict[str, str] = {}
62        add_anthropic_available_models_env(env, '[{"name":"model"}]')
63
64        self.assertEqual(env["ZED_ANTHROPIC_AVAILABLE_MODELS"], '[{"name":"model"}]')
65
66    def test_configure_modal_environment_sets_documented_secret_env(self) -> None:
67        args = argparse.Namespace(
68            app_name="app",
69            volume="volume",
70            modal_token_secret="modal-secret",
71            api_secret="llm-secret",
72        )
73
74        with patch.dict(os.environ, {}, clear=False):
75            common.configure_modal_environment(args)
76
77            self.assertEqual(os.environ["AGENT_EVALS_APP_NAME"], "app")
78            self.assertEqual(os.environ["AGENT_EVALS_VOLUME"], "volume")
79            self.assertEqual(
80                os.environ["AGENT_EVALS_MODAL_TOKEN_SECRET"], "modal-secret"
81            )
82            self.assertEqual(
83                os.environ["AGENT_EVALS_LLM_PROVIDERS_SECRET"], "llm-secret"
84            )
85
86
87class SafeExtractArchiveTests(unittest.TestCase):
88    def make_archive(
89        self, archive_path: Path, member_name: str, data: bytes = b"data"
90    ) -> None:
91        with tarfile.open(archive_path, "w:gz") as archive:
92            info = tarfile.TarInfo(member_name)
93            info.size = len(data)
94            archive.addfile(info, io.BytesIO(data))
95
96    def test_extracts_normal_members(self) -> None:
97        with tempfile.TemporaryDirectory() as temporary_directory:
98            root = Path(temporary_directory)
99            archive_path = root / "archive.tar.gz"
100            destination = root / "out"
101            destination.mkdir()
102            self.make_archive(archive_path, "job/result.json", b"{}")
103
104            with tarfile.open(archive_path, "r:gz") as archive:
105                common.safe_extract_archive(archive, destination)
106
107            self.assertEqual((destination / "job" / "result.json").read_text(), "{}")
108
109    def test_rejects_path_traversal_members(self) -> None:
110        with tempfile.TemporaryDirectory() as temporary_directory:
111            root = Path(temporary_directory)
112            archive_path = root / "archive.tar.gz"
113            destination = root / "out"
114            destination.mkdir()
115            self.make_archive(archive_path, "../evil.txt")
116
117            with tarfile.open(archive_path, "r:gz") as archive:
118                with self.assertRaises(ValueError):
119                    common.safe_extract_archive(archive, destination)
120
121    def test_rejects_links(self) -> None:
122        with tempfile.TemporaryDirectory() as temporary_directory:
123            root = Path(temporary_directory)
124            archive_path = root / "archive.tar.gz"
125            destination = root / "out"
126            destination.mkdir()
127            with tarfile.open(archive_path, "w:gz") as archive:
128                info = tarfile.TarInfo("link")
129                info.type = tarfile.SYMTYPE
130                info.linkname = "/tmp/target"
131                archive.addfile(info)
132
133            with tarfile.open(archive_path, "r:gz") as archive:
134                with self.assertRaises(ValueError):
135                    common.safe_extract_archive(archive, destination)
136
137
138if __name__ == "__main__":
139    unittest.main()
140