303 lines
14 KiB
Python
303 lines
14 KiB
Python
import json
|
|
import os
|
|
import subprocess
|
|
import sys
|
|
import tempfile
|
|
import unittest
|
|
from pathlib import Path
|
|
|
|
|
|
REPO_ROOT = Path(__file__).resolve().parents[1]
|
|
sys.path.insert(0, str(REPO_ROOT))
|
|
|
|
import host
|
|
|
|
|
|
class HostPureFunctionsTest(unittest.TestCase):
|
|
def test_save_state_merges_tasks_and_preserves_disk_tasks(self):
|
|
with tempfile.TemporaryDirectory() as tmp:
|
|
state_path = os.path.join(tmp, "host-state.json")
|
|
with open(state_path, "w", encoding="utf-8") as f:
|
|
json.dump({"tasks": {"disk-task": {"status": "running"}},
|
|
"automations": []}, f)
|
|
|
|
state = {"tasks": {"memory-task": {"status": "completed"}},
|
|
"automations": []}
|
|
host.save_state(state_path, state)
|
|
|
|
self.assertEqual(state["tasks"]["disk-task"], {"status": "running"})
|
|
self.assertEqual(state["tasks"]["memory-task"], {"status": "completed"})
|
|
loaded = host.load_state(state_path)
|
|
self.assertEqual(loaded, state)
|
|
|
|
def test_task_diff_includes_untracked_text_contents(self):
|
|
with tempfile.TemporaryDirectory() as tmp:
|
|
subprocess.run(["git", "init", "-q"], cwd=tmp, check=True)
|
|
Path(tmp, "tracked.txt").write_text("tracked\n", encoding="utf-8")
|
|
subprocess.run(["git", "add", "tracked.txt"], cwd=tmp, check=True)
|
|
subprocess.run(["git", "-c", "user.name=Test", "-c",
|
|
"user.email=test@example.com", "commit", "-qm",
|
|
"initial"], cwd=tmp, check=True)
|
|
Path(tmp, "new.txt").write_text("new file contents\n", encoding="utf-8")
|
|
|
|
result = host.task_diff(tmp)
|
|
|
|
self.assertIn("new.txt", result["untracked"].splitlines())
|
|
self.assertEqual(result["untracked_contents"]["new.txt"],
|
|
"new file contents\n")
|
|
|
|
def test_load_config_file_returns_empty_dict_for_missing_file(self):
|
|
with tempfile.TemporaryDirectory() as tmp:
|
|
missing = os.path.join(tmp, "missing.json")
|
|
self.assertEqual(host._load_config_file(missing), {})
|
|
|
|
|
|
def _init_git_repo(tmp):
|
|
subprocess.run(["git", "init", "-q"], cwd=tmp, check=True)
|
|
subprocess.run(["git", "config", "user.name", "Test"], cwd=tmp, check=True)
|
|
subprocess.run(["git", "config", "user.email", "test@example.com"],
|
|
cwd=tmp, check=True)
|
|
|
|
|
|
class TopoSortTest(unittest.TestCase):
|
|
def test_priority_and_dependency_order(self):
|
|
jobs = [
|
|
{"name": "low", "depends_on": [], "priority": 0},
|
|
{"name": "high", "depends_on": [], "priority": 5},
|
|
{"name": "child", "depends_on": ["low", "high"], "priority": 0},
|
|
]
|
|
order = [j["name"] for j in host.topo_sort_jobs(jobs)]
|
|
self.assertEqual(order[0], "high")
|
|
self.assertEqual(order[-1], "child")
|
|
|
|
def test_cycle_and_missing_dependency_raise(self):
|
|
with self.assertRaises(ValueError):
|
|
host.topo_sort_jobs([{"name": "a", "depends_on": ["b"]}])
|
|
with self.assertRaises(ValueError):
|
|
host.topo_sort_jobs([
|
|
{"name": "a", "depends_on": ["b"]},
|
|
{"name": "b", "depends_on": ["a"]},
|
|
])
|
|
|
|
|
|
class MergeVerifyTest(unittest.TestCase):
|
|
def _make_repo(self, tmp, break_test=False):
|
|
_init_git_repo(tmp)
|
|
Path(tmp, "mod.py").write_text("def f():\n return 1\n",
|
|
encoding="utf-8")
|
|
tests_dir = Path(tmp, "tests")
|
|
tests_dir.mkdir()
|
|
body = ("import unittest\nfrom mod import f\n\n"
|
|
"class T(unittest.TestCase):\n"
|
|
" def test_f(self):\n"
|
|
" self.assertEqual(f(), 1)\n")
|
|
(tests_dir / "test_mod.py").write_text(body, encoding="utf-8")
|
|
subprocess.run(["git", "add", "-A"], cwd=tmp, check=True)
|
|
subprocess.run(["git", "commit", "-qm", "initial"], cwd=tmp,
|
|
check=True)
|
|
# worktree 手动建：模拟 host 的 worktree/分支
|
|
subprocess.run(["git", "worktree", "add", os.path.join(tmp, "wt"),
|
|
"-b", "host/t1", "HEAD"], cwd=tmp, check=True)
|
|
if break_test:
|
|
Path(tmp, "wt", "mod.py").write_text(
|
|
"def f():\n return 2\n", encoding="utf-8")
|
|
else:
|
|
Path(tmp, "wt", "mod.py").write_text(
|
|
"def f():\n return 1 # comment\n", encoding="utf-8")
|
|
return os.path.join(tmp, "wt")
|
|
|
|
def test_merge_verify_pass_cleans_worktree(self):
|
|
with tempfile.TemporaryDirectory() as tmp:
|
|
wt = self._make_repo(tmp)
|
|
ok, msg = host.merge_task(tmp, wt, "host/t1", "t1",
|
|
verify_command=None)
|
|
self.assertTrue(ok, msg)
|
|
self.assertFalse(os.path.isdir(wt))
|
|
out = subprocess.run(["git", "branch", "--list", "host/t1"],
|
|
cwd=tmp, capture_output=True, text=True)
|
|
self.assertEqual(out.stdout.strip(), "")
|
|
|
|
def test_merge_verify_fail_rolls_back_and_keeps_worktree(self):
|
|
with tempfile.TemporaryDirectory() as tmp:
|
|
wt = self._make_repo(tmp, break_test=True)
|
|
pre = subprocess.run(["git", "rev-parse", "HEAD"], cwd=tmp,
|
|
capture_output=True, text=True).stdout.strip()
|
|
ok, msg = host.merge_task(tmp, wt, "host/t1", "t1",
|
|
verify_command=None)
|
|
self.assertFalse(ok)
|
|
self.assertIn("回滚", msg)
|
|
now = subprocess.run(["git", "rev-parse", "HEAD"], cwd=tmp,
|
|
capture_output=True, text=True).stdout.strip()
|
|
self.assertEqual(now, pre)
|
|
self.assertTrue(os.path.isdir(wt))
|
|
self.assertEqual(Path(tmp, "mod.py").read_text(),
|
|
"def f():\n return 1\n")
|
|
|
|
def test_merge_duplicate_add_auto_resolved_with_ours(self):
|
|
# 依赖任务场景：两个分支都在同一位置加了同一段 add2，
|
|
# 先合并第一个，再合并第二个时应自动解（-X ours），
|
|
# 且第二个独有的 mul2 仍在主仓里。
|
|
with tempfile.TemporaryDirectory() as tmp:
|
|
_init_git_repo(tmp)
|
|
Path(tmp, "mod.py").write_text("def f():\n return 1\n",
|
|
encoding="utf-8")
|
|
subprocess.run(["git", "add", "-A"], cwd=tmp, check=True)
|
|
subprocess.run(["git", "commit", "-qm", "initial"], cwd=tmp,
|
|
check=True)
|
|
# 分支1：加 add2
|
|
subprocess.run(["git", "worktree", "add",
|
|
os.path.join(tmp, "wt1"), "-b", "host/t1",
|
|
"HEAD"], cwd=tmp, check=True)
|
|
Path(tmp, "wt1", "mod.py").write_text(
|
|
"def f():\n return 1\n\ndef add2(a):\n return a + 2\n",
|
|
encoding="utf-8")
|
|
ok, _ = host.merge_task(tmp, os.path.join(tmp, "wt1"),
|
|
"host/t1", "t1", verify_command="")
|
|
self.assertTrue(ok)
|
|
# 分支2：从初始 HEAD 建，但拷了分支1的改动 + 自己的 mul2
|
|
subprocess.run(["git", "worktree", "add",
|
|
os.path.join(tmp, "wt2"), "-b", "host/t2",
|
|
"HEAD"], cwd=tmp, check=True)
|
|
Path(tmp, "wt2", "mod.py").write_text(
|
|
"def f():\n return 1\n\ndef add2(a):\n return a + 2\n\n"
|
|
"def mul2(a):\n return a * 2\n", encoding="utf-8")
|
|
ok, msg = host.merge_task(tmp, os.path.join(tmp, "wt2"),
|
|
"host/t2", "t2", verify_command="")
|
|
self.assertTrue(ok, msg)
|
|
content = Path(tmp, "mod.py").read_text()
|
|
self.assertIn("def add2", content)
|
|
self.assertIn("def mul2", content)
|
|
# add2 只出现一次，不是重复两份
|
|
self.assertEqual(content.count("def add2"), 1)
|
|
|
|
def test_merge_does_not_commit_pycache(self):
|
|
with tempfile.TemporaryDirectory() as tmp:
|
|
wt = self._make_repo(tmp)
|
|
# 模拟 agent 跑测试生成的缓存
|
|
pyc_dir = Path(wt, "__pycache__")
|
|
pyc_dir.mkdir(exist_ok=True)
|
|
(pyc_dir / "mod.cpython-312.pyc").write_bytes(b"fake-pyc")
|
|
ok, msg = host.merge_task(tmp, wt, "host/t1", "t1",
|
|
verify_command="")
|
|
self.assertTrue(ok, msg)
|
|
tracked = subprocess.run(["git", "ls-files"], cwd=tmp,
|
|
capture_output=True, text=True).stdout
|
|
self.assertNotIn("pyc", tracked)
|
|
self.assertNotIn("__pycache__", tracked)
|
|
|
|
|
|
class ReviewCommentTest(unittest.TestCase):
|
|
def test_comment_and_fix_prompt(self):
|
|
task = {"name": "t", "review_comments": []}
|
|
host.add_review_comment(task, "mod.py", 12, "加 docstring")
|
|
host.add_review_comment(task, "mod.py", 20, "处理空输入")
|
|
self.assertEqual(len(host.open_review_comments(task)), 2)
|
|
prompt = host.build_fix_prompt(task)
|
|
self.assertIn("mod.py:12: 加 docstring", prompt)
|
|
self.assertIn("mod.py:20: 处理空输入", prompt)
|
|
|
|
|
|
class RetryAndBudgetTest(unittest.TestCase):
|
|
class FakeServer:
|
|
def __init__(self, results):
|
|
self.results = list(results)
|
|
self.calls = 0
|
|
|
|
def run_turn(self, thread_id, prompt, model=None, effort=None,
|
|
sandbox=None):
|
|
self.calls += 1
|
|
return self.results.pop(0)
|
|
|
|
def test_retry_then_success(self):
|
|
task = {"name": "t", "thread_id": "x", "sandbox": None,
|
|
"max_retries": 2, "retry_count": 0, "turns": []}
|
|
state = {"tasks": {"t": task}, "automations": []}
|
|
srv = self.FakeServer([
|
|
{"status": "failed", "text": ""},
|
|
{"status": "completed", "text": "ok"},
|
|
])
|
|
result = host.run_task_turn_with_retry(srv, state, task, "do it",
|
|
"m", "low", cfg={})
|
|
self.assertEqual(result["status"], "completed")
|
|
self.assertEqual(task["retry_count"], 1)
|
|
self.assertEqual(srv.calls, 2)
|
|
|
|
def test_budget_blocks_before_turn(self):
|
|
task = {"name": "t", "thread_id": "x", "sandbox": None,
|
|
"max_retries": 0, "retry_count": 0, "turns": []}
|
|
today = __import__("time").strftime("%Y-%m-%d")
|
|
state = {"tasks": {"t": task}, "automations": [],
|
|
"usage": {"date": today, "turns": 5}}
|
|
srv = self.FakeServer([{"status": "completed", "text": "ok"}])
|
|
result = host.run_task_turn_with_retry(
|
|
srv, state, task, "do it", "m", "low",
|
|
cfg={"max_daily_turns": 5})
|
|
self.assertEqual(result["status"], "failed")
|
|
self.assertEqual(srv.calls, 0)
|
|
|
|
def test_notify_failure_writes_log_and_runs_command(self):
|
|
with tempfile.TemporaryDirectory() as tmp:
|
|
state_path = os.path.join(tmp, "host-state.json")
|
|
notify_log = os.path.join(tmp, "notify.log")
|
|
cfg = {"notify_command": f"echo {{task}} {{status}} >> {notify_log}"}
|
|
host.notify_failure(state_path, cfg, "t1",
|
|
{"status": "failed", "text": "boom"})
|
|
self.assertTrue(os.path.exists(
|
|
os.path.join(tmp, "host-run.log")))
|
|
self.assertIn("t1 failed", Path(notify_log).read_text())
|
|
|
|
|
|
class StateVersionTest(unittest.TestCase):
|
|
def test_version_increments_and_conflict_logged(self):
|
|
with tempfile.TemporaryDirectory() as tmp:
|
|
path = os.path.join(tmp, "host-state.json")
|
|
s1 = {"tasks": {"a": {"status": "completed"}}, "automations": []}
|
|
host.save_state(path, s1)
|
|
with open(path) as f:
|
|
self.assertEqual(json.load(f)["version"], 1)
|
|
# 模拟另一写者：直接改磁盘版本为 5，再保存 s1（加载版本是 1）
|
|
with open(path) as f:
|
|
disk = json.load(f)
|
|
disk["version"] = 5
|
|
disk["tasks"]["b"] = {"status": "completed"}
|
|
with open(path, "w") as f:
|
|
json.dump(disk, f)
|
|
host.save_state(path, s1)
|
|
with open(path) as f:
|
|
merged = json.load(f)
|
|
self.assertEqual(merged["version"], 6)
|
|
self.assertIn("b", merged["tasks"])
|
|
self.assertIn("a", merged["tasks"])
|
|
log = Path(tmp, "host-run.log").read_text()
|
|
self.assertIn("STATE-CONFLICT", log)
|
|
|
|
|
|
class ServerRestartTest(unittest.TestCase):
|
|
def test_kill_then_ensure_alive(self):
|
|
# 不跑模型：只验证 server 进程死了能重拉、list 能通。
|
|
# 用配对二进制（VM 本机补丁），隔离 CODEX_HOME。
|
|
import host as host_mod
|
|
with tempfile.TemporaryDirectory() as tmp:
|
|
codex_home = os.path.join(tmp, "codex-home")
|
|
srv = host_mod.Server(codex_home=codex_home, approval="deny")
|
|
try:
|
|
srv.initialize()
|
|
self.assertTrue(srv.is_alive())
|
|
srv.proc.kill()
|
|
srv.proc.wait(timeout=10)
|
|
self.assertFalse(srv.is_alive())
|
|
self.assertTrue(srv.ensure_alive())
|
|
self.assertTrue(srv.is_alive())
|
|
# list_threads 应能通（空或已有 thread 列表均可）
|
|
srv.list_threads()
|
|
finally:
|
|
try:
|
|
srv.close()
|
|
except Exception:
|
|
pass
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|
|
|