"""cloud worker 单测：认领顺序、取消、缺 key、失败写回、预算门。 运行（在 ~/workspace/codex-app-server 下）： python3 -m unittest discover -s cloud/tests -p "test_worker.py" -v 说明：只用注入的假 runner（假 CLI），不跑真实模型回合、不联网、 不登录任何服务；env 映射也是注入的普通 dict，里面的值只是 占位用的假名，绝不是真实密钥。 """ import json import os import sqlite3 import sys import tempfile import time import unittest from pathlib import Path from types import SimpleNamespace try: from cloud.cost import CostStore from cloud.models import ModelRouter from cloud.server import init_db from cloud.worker import Worker, resolve_env_key except ImportError: # discover -s cloud/tests 时把项目根目录补进 sys.path sys.path.insert(0, str(Path(__file__).resolve().parents[2])) from cloud.cost import CostStore from cloud.models import ModelRouter from cloud.server import init_db from cloud.worker import Worker, resolve_env_key def _insert_task(db_path, task_id, prompt="do work", status="queued", depends_on=None, priority=0, project_id="p1"): now = time.time() conn = sqlite3.connect(db_path) try: conn.execute( "INSERT INTO tasks(id, project_id, prompt, status," " idempotency_key, depends_on, priority, created_by," " created_at, updated_at, last_heartbeat)" " VALUES(?, ?, ?, ?, NULL, ?, ?, 'tester', ?, ?, NULL)", (task_id, project_id, prompt, status, json.dumps(depends_on or []), priority, now, now), ) conn.commit() finally: conn.close() def _status(db_path, task_id): conn = sqlite3.connect(db_path) try: row = conn.execute( "SELECT status FROM tasks WHERE id = ?", (task_id,)).fetchone() finally: conn.close() return row[0] if row else None def _completed_runner(cmd): """假 CLI：dispatch/run 一律成功，打印 host.py 风格的 JSON。""" if "dispatch" in cmd: name = cmd[cmd.index("--name") + 1] else: name = cmd[cmd.index("--task") + 1] return SimpleNamespace( returncode=0, stdout=json.dumps({"task": name, "status": "completed", "text": "done"}), stderr="") class WorkerTestCase(unittest.TestCase): def setUp(self): self.tmp = tempfile.TemporaryDirectory() self.db_path = os.path.join(self.tmp.name, "control.db") init_db(self.db_path) self.env = {"XCM_API_KEY": "test-placeholder"} self.calls = [] def tearDown(self): self.tmp.cleanup() def _worker(self, runner=None, **kwargs): kwargs.setdefault("env", self.env) return Worker(self.db_path, "/fake/host.py", runner=runner or _completed_runner, **kwargs) def _recording_runner(self, inner): def fake(cmd): self.calls.append(list(cmd)) return inner(cmd) return fake # -- 认领顺序：依赖拓扑 + 优先级 ------------------------------------ def test_claim_order_dependency_and_priority(self): _insert_task(self.db_path, "A", priority=0) _insert_task(self.db_path, "B", depends_on=["A"], priority=0) _insert_task(self.db_path, "C", priority=9) worker = self._worker(runner=self._recording_runner(_completed_runner)) results = worker.run_all() self.assertEqual([r["status"] for r in results], ["completed"] * 3) # 首波 {C, A} 按优先级降序：C 先于 A；B 依赖 A，只能最后。 dispatch_names = [c[c.index("--name") + 1] for c in self.calls if "dispatch" in c] self.assertEqual(dispatch_names, ["C", "A", "B"]) for task_id in ("A", "B", "C"): self.assertEqual(_status(self.db_path, task_id), "completed") def test_dependency_on_failed_task_not_dispatched(self): def fake(cmd): if "dispatch" in cmd and cmd[cmd.index("--name") + 1] == "A": return SimpleNamespace(returncode=1, stdout="boom", stderr="") return _completed_runner(cmd) _insert_task(self.db_path, "A") _insert_task(self.db_path, "B", depends_on=["A"]) worker = self._worker(runner=fake) results = worker.run_all() self.assertEqual(len(results), 1) self.assertEqual(results[0]["status"], "failed") self.assertEqual(_status(self.db_path, "A"), "failed") # B 的依赖未完成，始终不可派，保持 queued。 self.assertEqual(_status(self.db_path, "B"), "queued") # -- 取消 ------------------------------------------------------------ def test_cancel_queued_task_marks_interrupted(self): _insert_task(self.db_path, "A") worker = self._worker(runner=self._recording_runner( _completed_runner)) self.assertTrue(worker.cancel("A")) self.assertEqual(_status(self.db_path, "A"), "interrupted") # 已终态：再次取消返回 False；不存在的任务也是 False。 self.assertFalse(worker.cancel("A")) self.assertFalse(worker.cancel("NOPE")) # 已取消的任务不再被派发。 self.assertEqual(worker.run_all(), []) self.assertEqual(self.calls, []) def test_cancel_during_dispatch_keeps_interrupted(self): _insert_task(self.db_path, "A") worker = None def fake(cmd): # 执行中途从外部取消：worker 返回后不得用 completed 覆盖。 worker.cancel("A") return _completed_runner(cmd) worker = self._worker(runner=fake) result = worker.run_one() self.assertEqual(result["status"], "interrupted") self.assertEqual(_status(self.db_path, "A"), "interrupted") # -- 缺 key 门 -------------------------------------------------------- def test_missing_key_fails_clearly_and_does_not_run(self): _insert_task(self.db_path, "A") worker = self._worker( runner=self._recording_runner(_completed_runner), env={}) result = worker.run_one() self.assertTrue(result["dispatched"]) self.assertEqual(result["status"], "failed") self.assertEqual(result["error"], "Missing environment variable: XCM_API_KEY") self.assertEqual(_status(self.db_path, "A"), "failed") self.assertEqual(self.calls, []) # 缺 key 时绝不调 host.py def test_missing_key_empty_value_counts_as_missing(self): _insert_task(self.db_path, "A") worker = self._worker(env={"XCM_API_KEY": ""}) result = worker.run_one() self.assertEqual(result["error"], "Missing environment variable: XCM_API_KEY") def test_custom_env_key_name_from_arg_and_config(self): # 参数指定的名字出现在报错里（只含名字，不含值）。 _insert_task(self.db_path, "A") worker = self._worker(env={}, env_key="MY_RELAY_KEY") self.assertEqual(worker.run_one()["error"], "Missing environment variable: MY_RELAY_KEY") # codex 配置文件里的 env_key 行只用来取名字。 cfg = os.path.join(self.tmp.name, "config.toml") with open(cfg, "w", encoding="utf-8") as f: f.write('env_key = "CUSTOM_RELAY_KEY"\n') self.assertEqual(resolve_env_key(codex_config_path=cfg), "CUSTOM_RELAY_KEY") _insert_task(self.db_path, "B") worker2 = self._worker(env={}, codex_config_path=cfg) self.assertEqual(worker2.env_key, "CUSTOM_RELAY_KEY") def test_key_present_proceeds(self): _insert_task(self.db_path, "A") worker = self._worker() result = worker.run_one() self.assertEqual(result["status"], "completed") self.assertIsNone(result["error"]) # -- 执行失败写回 + 成本 ---------------------------------------------- def test_execution_failure_writes_back_failed_and_cost(self): def fake(cmd): return SimpleNamespace(returncode=1, stdout="boom", stderr="") _insert_task(self.db_path, "A", prompt="do work") worker = self._worker(runner=fake) result = worker.run_one() self.assertEqual(result["status"], "failed") self.assertIn("boom", result["error"]) self.assertEqual(_status(self.db_path, "A"), "failed") store = CostStore(self.db_path) try: report = store.report_by_task("A") row = store._conn.execute( "SELECT model, source, fallback_count FROM cost_records" " WHERE task_id = ?", ("A",)).fetchone() finally: store._conn.close() self.assertEqual(report["turns"], 1) self.assertGreater(report["estimated_tokens"], 0) self.assertEqual(report["actual_tokens"], 0) # 复用 cost.py 的记录结构：估算标 estimated，不冒充 actual。 self.assertEqual(row[0], "gpt-6.1-sol") self.assertEqual(row[1], "estimated") self.assertEqual(row[2], 0) # -- 预算门 ------------------------------------------------------------ def test_budget_exceeded_no_dispatch(self): _insert_task(self.db_path, "A") worker = self._worker( runner=self._recording_runner(_completed_runner), max_turns=0) result = worker.run_one() self.assertFalse(result["dispatched"]) self.assertEqual(result["reason"], "budget_exceeded") self.assertEqual(_status(self.db_path, "A"), "queued") self.assertEqual(self.calls, []) def test_cost_budget_exceeded_no_dispatch(self): _insert_task(self.db_path, "A") worker = self._worker(max_cost_usd=0.0) result = worker.run_one() self.assertFalse(result["dispatched"]) self.assertEqual(_status(self.db_path, "A"), "queued") # -- fallback 事件：只记、不自动重试 ----------------------------------- def test_timeout_fallback_event_recorded_not_retried(self): def fake(cmd): return SimpleNamespace( returncode=0, stdout=json.dumps({"task": "A", "status": "timeout", "text": ""}), stderr="") _insert_task(self.db_path, "A") router = ModelRouter(["gpt-6.1-sol", "gpt-6-sol"]) worker = self._worker(runner=self._recording_runner(fake), router=router) result = worker.run_one() self.assertEqual(result["status"], "failed") self.assertEqual(len(result["events"]), 1) event = result["events"][0] self.assertEqual(event["type"], "model_fallback") self.assertEqual(event["from_model"], "gpt-6.1-sol") self.assertEqual(event["to_model"], "gpt-6-sol") self.assertEqual(event["reason"], "timeout") # 只记事件不重试：假 CLI 只被调用一次。 self.assertEqual(len(self.calls), 1) store = CostStore(self.db_path) try: report = store.report_by_task("A") finally: store._conn.close() self.assertEqual(report["fallback_count"], 1) # -- dispatch 已存在时用 run 续跑 -------------------------------------- def test_dispatch_existing_falls_back_to_run(self): def fake(cmd): if "dispatch" in cmd: return SimpleNamespace( returncode=1, stdout="", stderr="Traceback: task 已存在：A，先换名或 discard 它") return _completed_runner(cmd) _insert_task(self.db_path, "A") worker = self._worker(runner=self._recording_runner(fake)) result = worker.run_one() self.assertEqual(result["status"], "completed") self.assertEqual(_status(self.db_path, "A"), "completed") kinds = ["dispatch" if "dispatch" in c else "run" for c in self.calls] self.assertEqual(kinds, ["dispatch", "run"]) run_cmd = [c for c in self.calls if "run" in c][0] self.assertEqual(run_cmd[run_cmd.index("--task") + 1], "A") if __name__ == "__main__": unittest.main()