"""覆盖率 78%->85% 第 1 步：cloud/server.py 主力缺口的真实行为测试。 只测 server.py 的助手函数与本地 DB 迁移/配额/限流逻辑：全部用 临时目录、临时 SQLite、临时 stub 脚本作输入，断言真实输出与 状态变化；不 mock 被测函数本身，不跑模型、不联网、不建生产任务。 """ import json import os import sqlite3 import sys import tempfile import time import unittest from decimal import Decimal from pathlib import Path try: from cloud import server as _srv from cloud.server import ( DbRateLimiter, RateLimiter, _apply_default_quota, _attachments_map, _backfill_default_quotas, _build_agent_chain, _chat_message_from_trace, _continue_message_to_dict, _default_executor, _default_quota_usd_micros, _deployed_commit, _ensure_continue_messages_table, _ensure_environments_table, _ensure_quotas_table, _ensure_task_checkpoints_table, _ensure_task_steps_table, _ensure_uploads_tables, _env_int, _environment_to_dict, _extract_result_list, _extract_thread, _extract_thread_list, _followup_to_dict, _global_admin_ids, _host_py_path, _json_default, _latest_ts, _agent_status_from_task, _local_worker_codex_home, _make_rate_limiter, _pick_latest_thread, _rate_limit_from_env, _read_agent_entries, _read_agent_entries_multi, _read_trace_entries, _read_trace_entries_multi, _resolve_shared_home_mcp, _run_host_cli, _safe_worktree_path, _task_tokens_map, _task_worktree_for_files, _tenant_codex_home, _tenant_safe_project_name, _tenant_trace_dir, _tenant_work_dir, _thread_worktree_for_task, _trace_dir_for, _trace_dirs_for_task, _trace_file_name, _tenants_root_for, _validate_env_branch, _worker_heartbeat_status, init_db, migrate_projects_table, migrate_tasks_table, sweep_rate_limit_hits, ) except ImportError: # discover -s cloud/tests 时把项目根目录补进 sys.path sys.path.insert(0, str(Path(__file__).resolve().parents[2])) from cloud import server as _srv from cloud.server import ( DbRateLimiter, RateLimiter, _apply_default_quota, _attachments_map, _backfill_default_quotas, _build_agent_chain, _chat_message_from_trace, _continue_message_to_dict, _default_executor, _default_quota_usd_micros, _deployed_commit, _ensure_continue_messages_table, _ensure_environments_table, _ensure_quotas_table, _ensure_task_checkpoints_table, _ensure_task_steps_table, _ensure_uploads_tables, _env_int, _environment_to_dict, _extract_result_list, _extract_thread, _extract_thread_list, _followup_to_dict, _global_admin_ids, _host_py_path, _json_default, _latest_ts, _agent_status_from_task, _local_worker_codex_home, _make_rate_limiter, _pick_latest_thread, _rate_limit_from_env, _read_agent_entries, _read_agent_entries_multi, _read_trace_entries, _read_trace_entries_multi, _resolve_shared_home_mcp, _run_host_cli, _safe_worktree_path, _task_tokens_map, _task_worktree_for_files, _tenant_codex_home, _tenant_safe_project_name, _tenant_trace_dir, _tenant_work_dir, _thread_worktree_for_task, _trace_dir_for, _trace_dirs_for_task, _trace_file_name, _tenants_root_for, _validate_env_branch, _worker_heartbeat_status, init_db, migrate_projects_table, migrate_tasks_table, sweep_rate_limit_hits, ) class _Env: """临时改环境变量，退出时原样恢复。""" def __init__(self, **kw): self.kw = kw self.old = {} def __enter__(self): for k, v in self.kw.items(): self.old[k] = os.environ.get(k) if v is None: os.environ.pop(k, None) else: os.environ[k] = v return self def __exit__(self, *exc): for k, v in self.old.items(): if v is None: os.environ.pop(k, None) else: os.environ[k] = v return False class _StubServer: def __init__(self, **attrs): for k, v in attrs.items(): setattr(self, k, v) class EnvIntTestCase(unittest.TestCase): def test_valid_and_defaults(self): with _Env(COV85_X="42"): self.assertEqual(_env_int("COV85_X", 7), 42) with _Env(COV85_X=None): self.assertEqual(_env_int("COV85_X", 7), 7) with _Env(COV85_X="not-a-number"): self.assertEqual(_env_int("COV85_X", 7), 7) with _Env(COV85_X="0"): self.assertEqual(_env_int("COV85_X", 7), 7) with _Env(COV85_X="-3"): self.assertEqual(_env_int("COV85_X", 7), 7) def test_rate_limit_from_env(self): with _Env(COV85_RL="5"): self.assertEqual(_rate_limit_from_env("COV85_RL", 120), 5) with _Env(COV85_RL="bad"): self.assertEqual(_rate_limit_from_env("COV85_RL", 120), 120) with _Env(COV85_RL="0"): self.assertEqual(_rate_limit_from_env("COV85_RL", 120), 120) class PathHelpersTestCase(unittest.TestCase): def setUp(self): self.tmp = tempfile.TemporaryDirectory() self.addCleanup(self.tmp.cleanup) def test_trace_file_name_sanitizes(self): self.assertEqual(_trace_file_name("abc/def:ghi"), "abc-def-ghi") self.assertEqual(_trace_file_name("a_b-c.d"), "a_b-c.d") def test_trace_dir_priority(self): srv = _StubServer(trace_dir="/tmp/from-attr") self.assertEqual(_trace_dir_for(srv), "/tmp/from-attr") with _Env(CODEX_TRACE_DIR="/tmp/from-env"): self.assertEqual(_trace_dir_for(_StubServer()), "/tmp/from-env") self.assertEqual(_trace_dir_for(None), "/tmp/from-env") def test_tenants_root_priority(self): srv = _StubServer(tenants_root="/tmp/tenants") self.assertEqual(_tenants_root_for(srv), "/tmp/tenants") with _Env(CODEX_TENANTS_ROOT="/tmp/env-tenants"): self.assertEqual(_tenants_root_for(_StubServer()), "/tmp/env-tenants") def test_tenant_trace_dir(self): srv = _StubServer(tenants_root="/tmp/tenants") self.assertEqual(_tenant_trace_dir(srv, "proj-1"), "/tmp/tenants/proj-1/work/host-traces") self.assertIsNone(_tenant_trace_dir(srv, "")) self.assertIsNone(_tenant_trace_dir(srv, "...")) self.assertIsNone(_tenant_trace_dir(srv, None)) def test_trace_dirs_for_task_tenant_first_no_dup(self): srv = _StubServer(tenants_root="/tmp/tenants", trace_dir="/tmp/shared") dirs = _trace_dirs_for_task(srv, {"project_id": "p1"}) self.assertEqual(dirs[0], "/tmp/tenants/p1/work/host-traces") self.assertEqual(dirs[1], "/tmp/shared") # 租户目录与共享目录相同时不重复。 srv2 = _StubServer(tenants_root="/tmp", trace_dir="/tmp/p1/work/host-traces") # tenants_root=/tmp 时 tenant dir 为 /tmp/p1/work/host-traces self.assertEqual(_trace_dirs_for_task(srv2, {"project_id": "p1"}), ["/tmp/p1/work/host-traces"]) def test_tenant_safe_project_name(self): self.assertEqual(_tenant_safe_project_name("a/b"), "a-b") self.assertIsNone(_tenant_safe_project_name("")) self.assertIsNone(_tenant_safe_project_name("...")) self.assertIsNone(_tenant_safe_project_name(None)) self.assertEqual(_tenant_codex_home( _StubServer(tenants_root="/tmp/t"), "p1"), "/tmp/t/p1/codex-home") self.assertIsNone(_tenant_codex_home( _StubServer(tenants_root="/tmp/t"), "...")) self.assertEqual(_tenant_work_dir( _StubServer(tenants_root="/tmp/t"), "p1"), "/tmp/t/p1/work") def test_thread_worktree_for_task(self): srv = _StubServer(host_worktrees_dir="/tmp/wt") self.assertEqual(_thread_worktree_for_task(srv, {"id": "t/1"}), "/tmp/wt/t-1") with _Env(CODEX_HOST_WORKTREES_DIR="/tmp/env-wt"): self.assertEqual(_thread_worktree_for_task( _StubServer(), {"id": "t1"}), "/tmp/env-wt/t1") def test_task_worktree_for_files(self): # cloud 任务 409。 path, err = _task_worktree_for_files( _StubServer(tenants_root="/tmp/t"), {"id": "t1", "executor": "cloud"}) self.assertIsNone(path) self.assertEqual(err[0], 409) # 异机 worktree 409。 import socket path, err = _task_worktree_for_files( _StubServer(tenants_root="/tmp/t"), {"id": "t1", "worktree_host": "definitely-not-this-host-xyz"}) self.assertIsNone(path) self.assertEqual(err[0], 409) self.assertIn("definitely-not-this-host-xyz", err[1]) # 租户目录存在时优先命中。 tenant = os.path.join(self.tmp.name, "p1", "work", "host-worktrees", "t1") os.makedirs(tenant) path, err = _task_worktree_for_files( _StubServer(tenants_root=self.tmp.name), {"id": "t1", "project_id": "p1"}) self.assertEqual(path, tenant) self.assertIsNone(err) # 两处都不存在时 404。 path, err = _task_worktree_for_files( _StubServer(tenants_root=self.tmp.name, host_worktrees_dir=os.path.join(self.tmp.name, "wt")), {"id": "nope", "project_id": "p9"}) self.assertIsNone(path) self.assertEqual(err[0], 404) def test_safe_worktree_path(self): base = os.path.join(self.tmp.name, "wt") os.makedirs(os.path.join(base, "sub")) self.assertEqual(_safe_worktree_path(base, ""), base) self.assertEqual(_safe_worktree_path(base, "sub"), os.path.realpath(os.path.join(base, "sub"))) self.assertIsNone(_safe_worktree_path(base, "/etc/passwd")) self.assertIsNone(_safe_worktree_path(base, "../escape")) self.assertIsNone(_safe_worktree_path(base, "sub/../../escape")) # 指向内部的符号链接放行，指出去的拦住。 os.symlink(os.path.join(base, "sub"), os.path.join(base, "link-in")) self.assertIsNotNone(_safe_worktree_path(base, "link-in")) os.symlink("/etc", os.path.join(base, "link-out")) self.assertIsNone(_safe_worktree_path(base, "link-out")) class ExtractHelpersTestCase(unittest.TestCase): def test_extract_thread_list(self): self.assertIsNone(_extract_thread_list("nope")) self.assertIsNone(_extract_thread_list({})) self.assertEqual(_extract_thread_list({"data": [1, 2]}), [1, 2]) self.assertEqual(_extract_thread_list({"threads": [3]}), [3]) self.assertEqual(_extract_thread_list({"result": {"data": [4]}}), [4]) self.assertEqual(_extract_thread_list({"result": {"threads": [5]}}), [5]) self.assertIsNone(_extract_thread_list({"result": {"x": 1}})) def test_extract_thread(self): self.assertIsNone(_extract_thread(None)) self.assertEqual(_extract_thread({"thread": {"id": "a"}}), {"id": "a"}) self.assertEqual(_extract_thread({"id": "b"}), {"id": "b"}) self.assertEqual(_extract_thread({"result": {"thread": {"id": "c"}}}), {"id": "c"}) self.assertEqual(_extract_thread({"result": {"id": "d"}}), {"id": "d"}) self.assertIsNone(_extract_thread({"result": {}})) def test_extract_result_list(self): self.assertIsNone(_extract_result_list([], ["data"])) self.assertEqual(_extract_result_list({"data": [1]}, ["data"]), [1]) self.assertEqual(_extract_result_list({"alt": [2]}, ["data", "alt"]), [2]) self.assertEqual(_extract_result_list({"result": {"data": [3]}}, ["data"]), [3]) self.assertIsNone(_extract_result_list({"result": {}}, ["data"])) def test_pick_latest_thread(self): threads = [ {"id": "old", "updatedAt": 10}, {"id": "new", "updatedAt": 20}, {"id": "arch", "updatedAt": 99, "archived": True}, {"id": "bad-ts", "updatedAt": "xx"}, "not-a-dict", ] self.assertEqual(_pick_latest_thread(threads)["id"], "new") self.assertEqual( _pick_latest_thread(threads, include_archived=True)["id"], "arch") # 只有 createdAt 时退到 createdAt；全缺失时按原顺序取后者。 self.assertEqual(_pick_latest_thread( [{"id": "a", "createdAt": 5}, {"id": "b", "createdAt": 9}])["id"], "b") self.assertEqual(_pick_latest_thread([{"id": "a"}, {"id": "b"}])["id"], "b") self.assertIsNone(_pick_latest_thread([])) self.assertIsNone(_pick_latest_thread( [{"id": "x", "archived": True}])) def test_latest_ts(self): self.assertEqual(_latest_ts([1, 5, 3]), 5) self.assertEqual(_latest_ts([None, "x", 2.5]), 2.5) self.assertIsNone(_latest_ts([])) self.assertIsNone(_latest_ts([None, "x"])) def test_agent_status_from_task(self): self.assertEqual(_agent_status_from_task("queued"), "running") self.assertEqual(_agent_status_from_task("merged"), "completed") self.assertEqual(_agent_status_from_task("merge-failed"), "failed") self.assertEqual(_agent_status_from_task("interrupted"), "interrupted") self.assertEqual(_agent_status_from_task("weird"), "unknown") class TraceReadTestCase(unittest.TestCase): def setUp(self): self.tmp = tempfile.TemporaryDirectory() self.addCleanup(self.tmp.cleanup) def _write(self, name, lines): path = os.path.join(self.tmp.name, name) with open(path, "w", encoding="utf-8") as f: for line in lines: f.write(line + "\n") return path def test_read_trace_entries(self): path = self._write("t.jsonl", [ '{"id": "1", "kind": "message", "text": "a"}', 'not json', '{"no-id": true}', '', '{"id": "1", "kind": "message", "text": "b"}', '{"id": "2", "kind": "exec"}', '["not", "a", "dict"]', ]) entries = _read_trace_entries(path) self.assertEqual([e["id"] for e in entries], ["1", "2"]) self.assertEqual(entries[0]["text"], "b") # 后到覆盖先到，保首次顺序 self.assertEqual(_read_trace_entries( os.path.join(self.tmp.name, "missing.jsonl")), []) def test_read_trace_entries_multi(self): p1 = self._write("a.jsonl", ['{"id": "1", "v": 1}']) p2 = self._write("b.jsonl", ['{"id": "1", "v": 2}', '{"id": "2", "v": 1}']) entries = _read_trace_entries_multi( [p1, p2, os.path.join(self.tmp.name, "missing.jsonl")]) self.assertEqual([e["id"] for e in entries], ["1", "2"]) self.assertEqual(entries[0]["v"], 2) def test_read_agent_entries_thread_scoped(self): path = self._write("a.jsonl", [ '{"id": "1", "thread": "t1", "v": 1}', '{"id": "1", "thread": "t2", "v": 2}', '{"id": "1", "thread": "t1", "v": 3}', ]) entries = _read_agent_entries(path) self.assertEqual(len(entries), 2) # 同 id 不同线程不互相覆盖 self.assertEqual(entries[0]["v"], 3) self.assertEqual(entries[1]["v"], 2) multi = _read_agent_entries_multi([path]) self.assertEqual(len(multi), 2) def test_build_agent_chain(self): # 主任务 + 模型子代理 + 委派子任务（done）三路。 parent = self._write("task1.jsonl", [ json.dumps({"id": "c1", "thread": "main-th", "kind": "message", "text": "hi", "ts": 10}), json.dumps({"id": "c2", "thread": "sub-th", "kind": "message", "text": "sub", "ts": 20}), json.dumps({"id": "s1", "kind": "agent-start", "name": "checker", "prompt": "check it", "trace": "child.jsonl", "ts": 30}), json.dumps({"id": "d1", "kind": "agent-done", "name": "checker", "status": "completed", "ts": 40}), ]) self._write("child.jsonl", [ json.dumps({"id": "cc1", "kind": "message", "text": "child", "ts": 35}), ]) agents = _build_agent_chain( {"id": "task1", "status": "running"}, [self.tmp.name]) by_id = {a["id"]: a for a in agents} self.assertEqual(by_id["main"]["status"], "running") self.assertEqual(len(by_id["main"]["calls"]), 1) self.assertEqual(by_id["sub-th"]["kind"], "model") self.assertEqual(by_id["sub-th"]["status"], "running") self.assertEqual(by_id["delegate:checker"]["status"], "completed") self.assertEqual(by_id["delegate:checker"]["prompt"], "check it") self.assertEqual(len(by_id["delegate:checker"]["calls"]), 1) self.assertEqual(by_id["delegate:checker"]["updated_ts"], 40) # 委派无 done 标记时进行中；父任务已完成时模型子代理未知。 agents2 = _build_agent_chain( {"id": "task1", "status": "completed"}, [self.tmp.name]) by_id2 = {a["id"]: a for a in agents2} self.assertEqual(by_id2["sub-th"]["status"], "unknown") class JsonAndRateLimitTestCase(unittest.TestCase): def test_json_default_decimal(self): self.assertEqual(_json_default(Decimal("42")), 42) self.assertIsInstance(_json_default(Decimal("42")), int) self.assertAlmostEqual(_json_default(Decimal("1.5")), 1.5) with self.assertRaises(TypeError): _json_default(object()) def test_rate_limiter_window_and_prune(self): now = [1000.0] lim = RateLimiter(per_minute=2, window_seconds=60, time_fn=lambda: now[0]) self.assertTrue(lim.allow("k")) self.assertTrue(lim.allow("k")) self.assertFalse(lim.allow("k")) # 满 2 次即挡 now[0] = 1061.0 # 窗口滑过，旧 hit 被修剪 self.assertTrue(lim.allow("k")) self.assertTrue(lim.allow("other")) # 不同 key 独立计数 def test_db_rate_limiter_and_sweep(self): tmp = tempfile.TemporaryDirectory() self.addCleanup(tmp.cleanup) db = os.path.join(tmp.name, "rl.db") lim = DbRateLimiter(db, per_minute=2, window_seconds=60, time_fn=lambda: 1000.0) self.assertTrue(lim.allow("k")) self.assertTrue(lim.allow("k")) self.assertFalse(lim.allow("k")) # sweep 删掉窗口外行：先以旧时间写一行，再 sweep。 conn = sqlite3.connect(db) conn.execute("INSERT INTO rate_limit_hits(hit_key, ts)" " VALUES('old', 1.0)") conn.commit() conn.close() self.assertTrue(sweep_rate_limit_hits(db, 60, now=2000.0)) conn = sqlite3.connect(db) rows = conn.execute( "SELECT hit_key FROM rate_limit_hits").fetchall() conn.close() self.assertNotIn(("old",), rows) def test_make_rate_limiter_sqlite(self): with _Env(CLOUD_DB_BACKEND="sqlite"): lim = _make_rate_limiter(":memory:") self.assertIsInstance(lim, RateLimiter) self.assertNotIsInstance(lim, DbRateLimiter) def test_start_sweeper_non_db_returns_none(self): from cloud.server import start_rate_limit_sweeper self.assertIsNone(start_rate_limit_sweeper( RateLimiter(), ":memory:")) class HostCliTestCase(unittest.TestCase): def setUp(self): self.tmp = tempfile.TemporaryDirectory() self.addCleanup(self.tmp.cleanup) def _stub(self, name, body): path = os.path.join(self.tmp.name, name) with open(path, "w", encoding="utf-8") as f: f.write(body) return path def test_run_host_cli_success_and_errors(self): ok_stub = self._stub("ok.py", "import json, sys\n" "print(json.dumps({'ok': 1}))\n") with _Env(CODEX_HOST_PY=ok_stub): self.assertEqual(_host_py_path(), ok_stub) ok, result = _run_host_cli(["list"], None) self.assertTrue(ok) self.assertEqual(result, {"ok": 1}) err_stub = self._stub("err.py", "import json, sys\n" "print(json.dumps({'error': 'boom'}))\n" "sys.exit(1)\n") with _Env(CODEX_HOST_PY=err_stub): ok, result = _run_host_cli(["list"], None) self.assertFalse(ok) self.assertEqual(result, "boom") # 解析成功但 exit 非零时走 exit 分支（先解析、后看返回码）。 plain_fail = self._stub("plain.py", "import json, sys\n" "print(json.dumps({'ok': 1}))\n" "sys.exit(3)\n") with _Env(CODEX_HOST_PY=plain_fail): ok, result = _run_host_cli(["list"], None) self.assertFalse(ok) self.assertIn("exit 3", result) garbage = self._stub("garbage.py", "print('not json')\n") with _Env(CODEX_HOST_PY=garbage): ok, result = _run_host_cli(["list"], None) self.assertFalse(ok) self.assertIn("not valid JSON", result) not_obj = self._stub("notobj.py", "import json\nprint(json.dumps([1, 2]))\n") with _Env(CODEX_HOST_PY=not_obj): ok, result = _run_host_cli(["list"], None) self.assertFalse(ok) self.assertIn("not a JSON object", result) err_ok = self._stub("errok.py", "import json\n" "print(json.dumps({'error': 'soft'}))\n") with _Env(CODEX_HOST_PY=err_ok): ok, result = _run_host_cli(["list"], None) self.assertFalse(ok) self.assertEqual(result, "soft") with _Env(CODEX_HOST_PY=os.path.join(self.tmp.name, "missing.py")): ok, result = _run_host_cli(["list"], None) self.assertFalse(ok) self.assertEqual(result, "host.py not found") sleeper = self._stub("sleep.py", "import time\ntime.sleep(30)\n") with _Env(CODEX_HOST_PY=sleeper): ok, result = _run_host_cli(["list"], None, timeout=1) self.assertFalse(ok) self.assertEqual(result, "host.py timed out") class LocalHomeTestCase(unittest.TestCase): def setUp(self): self.tmp = tempfile.TemporaryDirectory() self.addCleanup(self.tmp.cleanup) self.stub = os.path.join(self.tmp.name, "host.py") with open(self.stub, "w", encoding="utf-8") as f: f.write("# stub\n") def _write_cfg(self, content): path = os.path.join(self.tmp.name, "host-config-host.json") with open(path, "w", encoding="utf-8") as f: f.write(content) return path def test_local_worker_codex_home_branches(self): with _Env(CODEX_HOST_PY=self.stub): # 无配置文件 -> None self.assertIsNone(_local_worker_codex_home()) self._write_cfg('{"codex_home": "/abs/home"}') self.assertEqual(_local_worker_codex_home(), "/abs/home") self._write_cfg('{"codex_home": "rel/home"}') self.assertEqual(_local_worker_codex_home(), os.path.join(self.tmp.name, "rel/home")) self._write_cfg('{"codex_home": " "}') self.assertIsNone(_local_worker_codex_home()) self._write_cfg('{"other": 1}') self.assertIsNone(_local_worker_codex_home()) self._write_cfg('{bad json') self.assertIsNone(_local_worker_codex_home()) # 兜底：本地解析失败时走 tenants_root 父目录。 resolved = _resolve_shared_home_mcp( _StubServer(tenants_root="/tmp/tenants")) self.assertEqual(resolved, "/tmp/codex-home") def test_resolve_shared_home_prefers_local(self): with _Env(CODEX_HOST_PY=self.stub): self._write_cfg('{"codex_home": "/abs/home"}') self.assertEqual(_resolve_shared_home_mcp( _StubServer(tenants_root="/tmp/tenants")), "/abs/home") class DbHelpersTestCase(unittest.TestCase): def setUp(self): self.tmp = tempfile.TemporaryDirectory() self.addCleanup(self.tmp.cleanup) self.db_path = os.path.join(self.tmp.name, "c.db") init_db(self.db_path) def _conn(self): conn = sqlite3.connect(self.db_path) conn.row_factory = sqlite3.Row return conn def test_migrate_tables_idempotent(self): conn = self._conn() try: migrate_tasks_table(conn) migrate_projects_table(conn) migrate_tasks_table(conn) # 二次跑不丢列、不报错 cols = [r[1] for r in conn.execute( "PRAGMA table_info(tasks)").fetchall()] self.assertIn("executor", cols) self.assertIn("worktree_host", cols) finally: conn.close() def test_ensure_tables_and_maps(self): conn = self._conn() try: _ensure_uploads_tables(conn) _ensure_task_checkpoints_table(conn) _ensure_task_steps_table(conn) _ensure_continue_messages_table(conn) _ensure_environments_table(conn) _ensure_quotas_table(conn) conn.commit() conn.execute("INSERT INTO task_attachments" "(task_id, bundle_id, rel_path, size)" " VALUES('t1', 'b1', 'b.txt', 2)," " ('t1', 'b1', 'a.txt', 1)") self.assertEqual(_attachments_map(conn, ["t1", "t2"]), {"t1": ["a.txt", "b.txt"], "t2": []}) self.assertEqual(_attachments_map(conn, []), {}) # cost_records 表不存在时全计 0，不抛错。 self.assertEqual(_task_tokens_map(conn, ["t1"]), {"t1": 0}) self.assertEqual(_task_tokens_map(conn, []), {}) finally: conn.close() def test_followup_and_continue_dicts(self): conn = self._conn() try: from cloud.peer_messages import ensure_followups_table ensure_followups_table(conn) conn.execute("INSERT INTO task_followups(id, task_id, project_id," " body, mode, status, created_by, created_at," " updated_at) VALUES('f1', 't1', 'p1', 'hi'," " 'normal', 'queued', 'u1', 1.0, 2.0)") row = conn.execute( "SELECT * FROM task_followups WHERE id='f1'").fetchone() d = _followup_to_dict(row) self.assertEqual(d["body"], "hi") self.assertIsNone(d["sender_task_id"]) self.assertEqual(d["created_at"], 1.0) _ensure_continue_messages_table(conn) conn.execute("INSERT INTO task_continue_messages" "(id, task_id, project_id, body, created_by," " created_at) VALUES('c1', 't1', 'p1', 'go'," " 'u1', 3.0)") row = conn.execute("SELECT * FROM task_continue_messages" " WHERE id='c1'").fetchone() d = _continue_message_to_dict(row) self.assertEqual(d, {"id": "c1", "task_id": "t1", "project_id": "p1", "body": "go", "created_by": "u1", "created_at": 3.0}) finally: conn.close() def test_chat_message_from_trace_kinds(self): self.assertEqual(_chat_message_from_trace( {"id": "1", "kind": "message", "phase": "final", "text": "hello"})["text"], "[final] hello") self.assertEqual(_chat_message_from_trace( {"id": "2", "kind": "message", "text": "x"})["role"], "assistant") exec_msg = _chat_message_from_trace( {"id": "3", "kind": "exec", "command": "ls", "exitCode": 0, "durationMs": 12, "cwd": "/tmp"}) self.assertIn("ls", exec_msg["text"]) self.assertEqual(exec_msg["label"], "命令") file_msg = _chat_message_from_trace( {"id": "4", "kind": "file", "path": "a.py", "change": "edit", "status": "ok"}) self.assertEqual(file_msg["text"], "a.py edit ok") tool_msg = _chat_message_from_trace( {"id": "5", "kind": "tool", "server": "s", "tool": "t", "status": "ok"}) self.assertEqual(tool_msg["text"], "s/t ok") call_msg = _chat_message_from_trace( {"id": "6", "kind": "call", "name": "fn"}) self.assertEqual(call_msg["text"], "fn") self.assertEqual(call_msg["label"], "调用") self.assertEqual(_chat_message_from_trace( {"id": "7", "kind": "plan", "text": "p"})["role"], "status") self.assertTrue(_chat_message_from_trace( {"id": "8", "kind": "reasoning", "summary": "hmm"})["text"].startswith("思考：")) self.assertIn("子代理开始", _chat_message_from_trace( {"id": "9", "kind": "agent-start", "name": "n", "prompt": "do"})["text"]) self.assertIn("子代理结束", _chat_message_from_trace( {"id": "10", "kind": "agent-done", "name": "n", "status": "completed"})["text"]) unknown = _chat_message_from_trace( {"id": "11", "kind": "weird", "x": 1}) self.assertIn("weird", unknown["text"]) def test_default_executor_and_quota(self): with _Env(CLOUD_DEFAULT_EXECUTOR=None): self.assertEqual(_default_executor(), "docker") with _Env(CLOUD_DEFAULT_EXECUTOR=""): self.assertIsNone(_default_executor()) with _Env(CLOUD_DEFAULT_EXECUTOR="local"): self.assertEqual(_default_executor(), "local") with _Env(CLOUD_DEFAULT_EXECUTOR="bogus"): self.assertEqual(_default_executor(), "docker") with _Env(CLOUD_DEFAULT_QUOTA_USD_MICROS="500"): self.assertEqual(_default_quota_usd_micros(), 500) with _Env(CLOUD_DEFAULT_QUOTA_USD_MICROS="bad"): self.assertEqual(_default_quota_usd_micros(), 10000000) with _Env(CLOUD_DEFAULT_QUOTA_USD_MICROS="-5"): self.assertEqual(_default_quota_usd_micros(), 0) def test_apply_and_backfill_quota(self): conn = self._conn() try: self.assertTrue(_apply_default_quota(conn, "u1")) self.assertFalse(_apply_default_quota(conn, "u1")) # 幂等不动 row = conn.execute("SELECT quota_usd_micros FROM user_quotas" " WHERE user_id='u1'").fetchone() self.assertEqual(row[0], 10000000) # backfill：users 里有、无配额行的补写；已有行不动。 # users 表由 AuthStore 建（init_db 只建 server 自己的表）。 from cloud.auth import AuthStore auth = AuthStore(self.db_path) u2_id, _ = auth.create_user("u2") u3_id, _ = auth.create_user("u3") conn.execute("INSERT INTO user_quotas(user_id," " quota_usd_micros, updated_at)" " VALUES(?, NULL, 1.0)", (u3_id,)) conn.commit() # u2 无行应被补写 1 行，u3 的 NULL 行不动；二次跑为 0。 self.assertEqual(_backfill_default_quotas(conn), 1) self.assertEqual(_backfill_default_quotas(conn), 0) row = conn.execute("SELECT quota_usd_micros FROM user_quotas" " WHERE user_id=?", (u2_id,)).fetchone() self.assertEqual(row[0], 10000000) row = conn.execute("SELECT quota_usd_micros FROM user_quotas" " WHERE user_id=?", (u3_id,)).fetchone() self.assertIsNone(row[0]) # NULL 不限额行不被覆盖 finally: conn.close() def test_global_admin_ids(self): with _Env(CODEX_GLOBAL_ADMINS="a, b,, c "): self.assertEqual(_global_admin_ids(), {"a", "b", "c"}) with _Env(CODEX_GLOBAL_ADMINS=None): self.assertEqual(_global_admin_ids(), set()) def test_validate_env_branch(self): self.assertIsNone(_validate_env_branch(None)) self.assertIsNone(_validate_env_branch(" ")) self.assertEqual(_validate_env_branch(" feat/x-1 "), "feat/x-1") for bad in ["-lead", "a..b", "x.lock", "a b", "a?b", "x" * 201]: with self.assertRaises(ValueError, msg=bad): _validate_env_branch(bad) with self.assertRaises(ValueError): _validate_env_branch(123) def test_environment_to_dict(self): conn = self._conn() try: _ensure_environments_table(conn) conn.execute("INSERT INTO environments(id, project_id, name," " branch, executor, status, error_text, repo_url," " created_by, created_at, updated_at)" " VALUES('e1', 'p1', 'base', 'main', 'docker'," " 'ready', NULL, 'https://example.invalid/r'," " 'u1', 1.0, 2.0)") row = conn.execute( "SELECT * FROM environments WHERE id='e1'").fetchone() d = _environment_to_dict(row) self.assertEqual(d["branch"], "main") self.assertEqual(d["executor"], "docker") self.assertIsNone(d["error_text"]) finally: conn.close() def test_worker_heartbeat_status(self): tmp = tempfile.TemporaryDirectory() self.addCleanup(tmp.cleanup) hb = os.path.join(tmp.name, "hb.json") with open(hb, "w", encoding="utf-8") as f: json.dump({"last_seen": time.time()}, f) srv = _StubServer(worker_heartbeat_file=hb, db_path=self.db_path) # 本库无 workers 行时回落文件心跳：新鲜文件 -> ok。 age, ok = _worker_heartbeat_status(srv) self.assertTrue(ok) self.assertIsNotNone(age) # 坏 JSON / 缺失文件 -> (None, False)，不抛错。 bad = os.path.join(tmp.name, "bad.json") with open(bad, "w", encoding="utf-8") as f: f.write("{oops") age, ok = _worker_heartbeat_status( _StubServer(worker_heartbeat_file=bad, db_path=self.db_path)) self.assertEqual((age, ok), (None, False)) age, ok = _worker_heartbeat_status(_StubServer( worker_heartbeat_file=os.path.join(tmp.name, "nope.json"), db_path=self.db_path)) self.assertEqual((age, ok), (None, False)) # workers 表有新鲜行时走注册表分支。 conn = sqlite3.connect(self.db_path) conn.execute("CREATE TABLE IF NOT EXISTS workers(" "worker_id TEXT PRIMARY KEY, last_seen REAL)") conn.execute("INSERT INTO workers(worker_id, last_seen)" " VALUES('w1', ?)", (time.time(),)) conn.commit() conn.close() age, ok = _worker_heartbeat_status(srv) self.assertTrue(ok) def test_deployed_commit_missing_returns_none(self): # 仓库根无 DEPLOYED_COMMIT 时返回 None（本 worktree 即如此， # 若未来部署产物落盘则此断言需随产物调整——此处只断言 # 类型：None 或非空字符串，且与文件内容一致）。 result = _deployed_commit() self.assertTrue(result is None or isinstance(result, str)) if __name__ == "__main__": unittest.main()