"""生产级 worker 超时 / 重试 / 容器清理兜底 + Postgres 备份单测。 只用标准库与本地 fake runner / 真实 sleep 进程，不跑模型、不联网； Postgres 备份用可注入的 fake pg_dump runner，不调真实 pg_dump， 密码只经 PGPASSWORD 环境变量传给子进程（单测里不放真实密码）。 """ import os import sqlite3 import subprocess import sys import tempfile import unittest from pathlib import Path from types import SimpleNamespace from unittest import mock try: from cloud import backup as backup_mod from cloud.auth import AuthStore from cloud.sandbox import DockerExecutor from cloud.worker import Worker except ImportError: sys.path.insert(0, str(Path(__file__).resolve().parents[2])) from cloud import backup as backup_mod from cloud.auth import AuthStore from cloud.sandbox import DockerExecutor from cloud.worker import Worker class _FakeRunner: """同步 fake 执行器 runner：按 results 顺序返回结果。 LocalExecutor 的 runner 收的是命令列表、返回 SimpleNamespace （returncode/stdout/stderr），与真实 subprocess 同款接口； 这里按顺序返回，不按任务 id 查表。 """ def __init__(self, results): self.results = list(results) self.calls = [] def __call__(self, cmd): self.calls.append(list(cmd)) spec = self.results[min(len(self.calls) - 1, len(self.results) - 1)] if isinstance(spec, Exception): raise spec return SimpleNamespace( returncode=0 if spec[0] == "completed" else 1, stdout=spec[1], stderr="") def _make_worker(tmp, runner=None, process_factory=None, task_timeout=1800, max_retries=1, retry_backoff_seconds=0): db_path = os.path.join(tmp, "control.db") # worker 本身不查项目成员关系：init_db + 直接插任务即可， # 不用走 AuthStore（其 create_project 只写成员表、不返项目）。 from cloud.server import init_db init_db(db_path) project_id = "proj" task_id = "t1" conn = sqlite3.connect(db_path) try: conn.execute( "INSERT INTO tasks(id, project_id, prompt, status," " depends_on, priority, created_by, created_at, updated_at," " last_heartbeat) VALUES(?, ?, ?, 'queued', '[]', 0," " 'system', 0, 0, NULL)", (task_id, project_id, "do it")) conn.commit() finally: conn.close() worker = Worker(db_path, "host.py", config_path=None, repo=None, runner=runner, env={"XCM_API_KEY": "test-placeholder"}, env_key=None, process_factory=process_factory, task_timeout=task_timeout, max_retries=max_retries, retry_backoff_seconds=retry_backoff_seconds) task = {"id": task_id, "project_id": project_id, "prompt": "do it", "depends_on": "[]", "priority": 0} return worker, db_path, task def _task_row(db_path): conn = sqlite3.connect(db_path) try: return conn.execute( "SELECT status, error_text FROM tasks WHERE id='t1'").fetchone() finally: conn.close() def _audit_actions(db_path): conn = sqlite3.connect(db_path) try: return [r[0] for r in conn.execute( "SELECT action FROM audit_log").fetchall()] finally: conn.close() class WorkerTimeoutTestCase(unittest.TestCase): def test_popen_timeout_kills_process_group_and_fails(self): tmp = tempfile.TemporaryDirectory() def factory(task): return subprocess.Popen( ["sleep", "30"], stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True, start_new_session=True) worker, db_path, task = _make_worker( tmp.name, process_factory=factory, task_timeout=0.2, max_retries=0, retry_backoff_seconds=0) result = worker.run_one(task) self.assertEqual(result["status"], "failed") # 超时是失败，不是取消：库里必须是 failed，绝不能是 interrupted。 status, error_text = _task_row(db_path) self.assertEqual(status, "failed") self.assertIn("timeout", error_text) self.assertIn("worker_timeout", _audit_actions(db_path)) tmp.cleanup() def test_task_timeout_zero_means_unlimited(self): # 对齐 host turn_timeout：0 归一化为 None（不限时）， # 而不是在 popen 轮询第一轮就立即超时。sleep 0.5 的 # 任务在 task_timeout=0 下应正常完成，库里是 completed。 tmp = tempfile.TemporaryDirectory() def factory(task): return subprocess.Popen( ["sleep", "0.5"], stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True, start_new_session=True) worker, db_path, task = _make_worker( tmp.name, process_factory=factory, task_timeout=0, max_retries=0, retry_backoff_seconds=0) self.assertIsNone(worker.task_timeout) result = worker.run_one(task) self.assertEqual(result["status"], "completed") status, _ = _task_row(db_path) self.assertEqual(status, "completed") tmp.cleanup() def test_docker_timeout_converts_to_failed_and_audit(self): tmp = tempfile.TemporaryDirectory() def docker_runner(cmd, timeout=None): if cmd[1] in ("kill", "rm"): return SimpleNamespace(returncode=0, stdout="", stderr="") raise subprocess.TimeoutExpired(cmd, timeout) worker, db_path, task = _make_worker( tmp.name, runner=_FakeRunner([]), max_retries=0, retry_backoff_seconds=0) # 测试占位 key（非真实密钥），只为通过执行器的“key 存在性” # 前置检查，让 runner 走到 TimeoutExpired 分支。 worker.executor = DockerExecutor( runner=docker_runner, env={"XCM_API_KEY": "test"}, timeout=5) # VM 上无 docker 二进制：只 mock 存在性检查（_check_docker）， # 真正的 run 仍走 fake runner 的 TimeoutExpired 分支。 with mock.patch("cloud.sandbox.shutil.which", return_value="/usr/bin/docker"): result = worker.run_one(task) self.assertEqual(result["status"], "failed") status, error_text = _task_row(db_path) self.assertEqual(status, "failed") self.assertIn("timeout", error_text) self.assertIn("worker_timeout", _audit_actions(db_path)) tmp.cleanup() class WorkerRetryTestCase(unittest.TestCase): def test_generic_failure_retries_once_then_failed(self): tmp = tempfile.TemporaryDirectory() runner = _FakeRunner([("failed", "boom"), ("failed", "boom")]) worker, db_path, task = _make_worker( tmp.name, runner=runner, max_retries=1, retry_backoff_seconds=0) result = worker.run_one(task) self.assertEqual(result["status"], "failed") self.assertEqual(len(runner.calls), 2) status, error_text = _task_row(db_path) self.assertEqual(status, "failed") self.assertIn("[retries=1]", error_text) self.assertIn("boom", error_text) self.assertIn("worker_retry", _audit_actions(db_path)) tmp.cleanup() def test_retry_success_recovers(self): tmp = tempfile.TemporaryDirectory() runner = _FakeRunner([("failed", "boom"), ("completed", "ok")]) worker, db_path, task = _make_worker( tmp.name, runner=runner, max_retries=1, retry_backoff_seconds=0) result = worker.run_one(task) self.assertEqual(result["status"], "completed") self.assertEqual(len(runner.calls), 2) status, _ = _task_row(db_path) self.assertEqual(status, "completed") tmp.cleanup() def test_max_retries_zero_means_no_retry(self): tmp = tempfile.TemporaryDirectory() runner = _FakeRunner([("failed", "boom")]) worker, db_path, task = _make_worker( tmp.name, runner=runner, max_retries=0, retry_backoff_seconds=0) result = worker.run_one(task) self.assertEqual(result["status"], "failed") self.assertEqual(len(runner.calls), 1) _, error_text = _task_row(db_path) self.assertNotIn("[retries=", error_text) tmp.cleanup() class ContainerCleanupTestCase(unittest.TestCase): def _docker_executor(self, ps_stdout=None, fail_ps=False): calls = [] def runner(cmd, timeout=None): calls.append(list(cmd)) if cmd[1] == "ps": if fail_ps: raise RuntimeError("docker down") return SimpleNamespace( returncode=0, stdout=ps_stdout or "", stderr="") return SimpleNamespace(returncode=0, stdout="", stderr="") executor = DockerExecutor(runner=runner) return executor, calls def test_cleanup_only_task_containers_never_infra(self): # 回归：2026-10-26 实测，按 codex- 前缀全删会把数据库容器 # codex-postgres 当成残留任务容器杀掉，导致控制面连不上 # Postgres。清理只能删 codex-<32 位 hex 任务 id>。 executor, calls = self._docker_executor( ps_stdout="codex-postgres\n" "codex-3689cf40d671492992c6bfedd18b004c\n" "codex-abc\nother\ncodex-app-server\n") cleaned = executor.cleanup_stale_containers() self.assertEqual( cleaned, ["codex-3689cf40d671492992c6bfedd18b004c"]) flat = [c for call in calls for c in call] self.assertNotIn("codex-postgres", flat) self.assertNotIn("codex-abc", flat) self.assertNotIn("other", flat) # kill + rm 各调一次（仅那个任务容器）。 self.assertEqual( sum(1 for call in calls if call[1] == "kill"), 1) self.assertEqual( sum(1 for call in calls if call[1] == "rm"), 1) def test_cleanup_failure_does_not_crash_worker(self): executor, _ = self._docker_executor(fail_ps=True) self.assertEqual(executor.cleanup_stale_containers(), []) tmp = tempfile.TemporaryDirectory() worker, _, _ = _make_worker(tmp.name, runner=_FakeRunner([])) worker.executor = executor # worker 层兜底同样不抛异常。 self.assertEqual(worker.cleanup_stale_containers(), []) tmp.cleanup() def test_list_containers_unparsable_or_failed_returns_empty(self): executor, _ = self._docker_executor(ps_stdout="") self.assertEqual(executor.list_containers(), []) class PostgresBackupTestCase(unittest.TestCase): def test_backup_postgres_fake_runner_success(self): tmp = tempfile.TemporaryDirectory() created = {} def fake_runner(cmd, env=None): self.assertNotIn("secret-pw", " ".join(cmd)) self.assertEqual(env.get("PGPASSWORD"), "secret-pw") dest = cmd[cmd.index("-f") + 1] Path(dest).write_text("-- dummy dump\n") created["cmd"] = cmd return SimpleNamespace(returncode=0, stdout="", stderr="") env = {"CLOUD_DB_NAME": "codexdb", "CLOUD_DB_USER": "codexsvc", "CLOUD_DB_PASSWORD": "secret-pw"} with mock.patch.dict(os.environ, env): dest = backup_mod.backup_postgres( tmp.name, runner=fake_runner) self.assertTrue(os.path.exists(dest)) self.assertGreater(os.path.getsize(dest), 0) self.assertNotIn("secret-pw", dest) self.assertIn("pg_dump", created["cmd"][0]) tmp.cleanup() def test_backup_postgres_failure_raises_clear_error(self): tmp = tempfile.TemporaryDirectory() def failing_runner(cmd, env=None): return SimpleNamespace( returncode=1, stdout="", stderr="pg_dump: boom") env = {"CLOUD_DB_NAME": "codexdb", "CLOUD_DB_USER": "codexsvc"} with mock.patch.dict(os.environ, env): with self.assertRaises(RuntimeError) as ctx: backup_mod.backup_postgres(tmp.name, runner=failing_runner) self.assertIn("pg_dump", str(ctx.exception)) tmp.cleanup() def test_backup_postgres_missing_binary_raises(self): tmp = tempfile.TemporaryDirectory() env = {"CLOUD_DB_NAME": "codexdb"} with mock.patch.dict(os.environ, env): with self.assertRaises(RuntimeError): backup_mod.backup_postgres( tmp.name, pg_dump_cmd="nonexistent-pg-dump-xyz") tmp.cleanup() if __name__ == "__main__": unittest.main()