Files
codex-app-server/cloud/tests/test_coverage85_server.py
host 488aedbcad code-to10 v6 A14 batch2: test E402/F401 cleanup plus touched-file residual E,F cleanup
(cherry picked from commit 1db02545f74c96b3b92be4fc32954659cec1f814)

Reconcile-conflict: A14b E402 imports referencing cloud.handler_core (does not exist here; core methods live in server.py) — kept HEAD imports with incoming's noqa: E402 markers, dropped the unused core_mod import (verified unused elsewhere in both files).
2026-11-03 02:49:26 +00:00

856 lines
36 KiB
Python

"""覆盖率 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.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.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。
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)三路。
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()