coverage85-step1: server/worker/host 主力缺口真实行为测试(+113 项)
只补 cloud/server.py、cloud/worker.py、host.py 的真实行为测试, 不动生产代码: - server:助手函数(路径/租户/提取/链路读/配额/限流/心跳/ host CLI stub 脚本)40 项 - worker:纯助手 + Worker DB/认领/容量/配额/跟进/租户/协调员 孤儿/回退 token,真实 SQLite 与真实子进程 36 项 - host:假 app-server 脚本走真实 JSON-RPC 协议,run_turn 成功/超时/中断/失败/历史错误回捞 + 状态/git/PR 助手 37 项 不含 mock 被测代码、无 skip、无 no-coverage、无空循环刷行。
This commit is contained in:
1 parent
f39ed3654e
commit
81a58daca2
3 files changed
+2677
No files matched your search
@@ -0,0 +1,858 @@
|
||||
"""覆盖率 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()
|
||||
@@ -0,0 +1,705 @@
|
||||
"""覆盖率 78%->85% 第 1 步:cloud/worker.py 主力缺口的真实行为测试。
|
||||
|
||||
只用临时 SQLite(server.init_db 建表)、临时目录与真实子进程:
|
||||
断言真实 DB 状态变化与返回值;不 mock 被测函数本身,不跑模型、
|
||||
不联网、不建生产任务。网关 sub2api 分支需要真实开通,刻意不碰。
|
||||
"""
|
||||
|
||||
import json
|
||||
import os
|
||||
import sqlite3
|
||||
import subprocess
|
||||
import sys
|
||||
import tempfile
|
||||
import time
|
||||
import unittest
|
||||
from pathlib import Path
|
||||
|
||||
try:
|
||||
from cloud.server import init_db
|
||||
from cloud.worker import (
|
||||
Worker,
|
||||
_code_signature,
|
||||
_default_quota_usd_micros,
|
||||
_merge_env,
|
||||
_missing_key_error,
|
||||
_resolve_model,
|
||||
_task_from_row,
|
||||
_worker_host_name,
|
||||
resolve_env_key,
|
||||
)
|
||||
except ImportError: # discover -s cloud/tests 时把项目根目录补进 sys.path
|
||||
sys.path.insert(0, str(Path(__file__).resolve().parents[2]))
|
||||
from cloud.server import init_db
|
||||
from cloud.worker import (
|
||||
Worker,
|
||||
_code_signature,
|
||||
_default_quota_usd_micros,
|
||||
_merge_env,
|
||||
_missing_key_error,
|
||||
_resolve_model,
|
||||
_task_from_row,
|
||||
_worker_host_name,
|
||||
resolve_env_key,
|
||||
)
|
||||
|
||||
|
||||
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
|
||||
|
||||
|
||||
def _insert_task(db_path, task_id, status="queued", **kw):
|
||||
conn = sqlite3.connect(db_path)
|
||||
try:
|
||||
cols = {"id": task_id, "project_id": kw.pop("project_id", "proj"),
|
||||
"prompt": kw.pop("prompt", "do it"), "status": status,
|
||||
"created_by": kw.pop("created_by", "owner1"),
|
||||
"created_at": kw.pop("created_at", 1000.0),
|
||||
"updated_at": kw.pop("updated_at", 1000.0)}
|
||||
cols.update(kw)
|
||||
names = ", ".join(cols)
|
||||
placeholders = ", ".join("?" * len(cols))
|
||||
conn.execute("INSERT INTO tasks(%s) VALUES(%s)"
|
||||
% (names, placeholders), tuple(cols.values()))
|
||||
conn.commit()
|
||||
finally:
|
||||
conn.close()
|
||||
|
||||
|
||||
class PureHelpersTestCase(unittest.TestCase):
|
||||
def test_worker_host_name(self):
|
||||
class Exec:
|
||||
def _host_name(self, task):
|
||||
return "host-" + task["id"]
|
||||
|
||||
self.assertEqual(_worker_host_name(Exec(), {"id": "t1"}), "host-t1")
|
||||
self.assertEqual(_worker_host_name(None, {"id": "t2"}), "t2")
|
||||
|
||||
class Raising:
|
||||
def _host_name(self, task):
|
||||
raise RuntimeError("boom")
|
||||
|
||||
self.assertEqual(_worker_host_name(Raising(), {"id": "t3"}), "t3")
|
||||
self.assertEqual(_worker_host_name(object(), {"id": "t4"}), "t4")
|
||||
|
||||
def test_default_quota(self):
|
||||
with _Env(CLOUD_DEFAULT_QUOTA_USD_MICROS="250"):
|
||||
self.assertEqual(_default_quota_usd_micros(), 250)
|
||||
with _Env(CLOUD_DEFAULT_QUOTA_USD_MICROS="bad"):
|
||||
self.assertEqual(_default_quota_usd_micros(), 10000000)
|
||||
with _Env(CLOUD_DEFAULT_QUOTA_USD_MICROS="-9"):
|
||||
self.assertEqual(_default_quota_usd_micros(), 0)
|
||||
|
||||
def test_code_signature_real_and_fake_host(self):
|
||||
sig = _code_signature("/fake/host.py")
|
||||
self.assertIsInstance(sig, tuple)
|
||||
self.assertEqual(len(sig), 6) # worker/sandbox/coordinator/cost/
|
||||
# server/host 共 6 个文件,全部取 (mtime_ns, size)。
|
||||
for mtime_ns, size in sig:
|
||||
self.assertIsInstance(mtime_ns, int)
|
||||
self.assertGreater(size, 0)
|
||||
sig2 = _code_signature()
|
||||
self.assertEqual(sig, sig2)
|
||||
|
||||
def test_missing_key_error(self):
|
||||
self.assertEqual(_missing_key_error("XCM_API_KEY"),
|
||||
"Missing environment variable: XCM_API_KEY")
|
||||
|
||||
def test_resolve_env_key(self):
|
||||
self.assertEqual(resolve_env_key("EXPLICIT"), "EXPLICIT")
|
||||
tmp = tempfile.TemporaryDirectory()
|
||||
self.addCleanup(tmp.cleanup)
|
||||
cfg = os.path.join(tmp.name, "config.toml")
|
||||
with open(cfg, "w", encoding="utf-8") as f:
|
||||
f.write('other = 1\nenv_key = "FROM_CFG"\n')
|
||||
self.assertEqual(resolve_env_key(None, cfg), "FROM_CFG")
|
||||
self.assertEqual(resolve_env_key(
|
||||
None, os.path.join(tmp.name, "missing.toml")),
|
||||
"XCM_API_KEY")
|
||||
with open(cfg, "w", encoding="utf-8") as f:
|
||||
f.write("no key here\n")
|
||||
self.assertEqual(resolve_env_key(None, cfg), "XCM_API_KEY")
|
||||
|
||||
def test_merge_env(self):
|
||||
self.assertEqual(_merge_env(None, {"a": 1}), {"a": 1})
|
||||
self.assertEqual(_merge_env({"a": 1}, None), {"a": 1})
|
||||
self.assertIsNone(_merge_env(None, None))
|
||||
merged = _merge_env({"a": 1, "b": 2}, {"b": 3, "c": 4})
|
||||
self.assertEqual(merged, {"a": 1, "b": 3, "c": 4})
|
||||
|
||||
def test_resolve_model(self):
|
||||
self.assertEqual(_resolve_model("m-explicit"), "m-explicit")
|
||||
tmp = tempfile.TemporaryDirectory()
|
||||
self.addCleanup(tmp.cleanup)
|
||||
cfg = os.path.join(tmp.name, "cfg.json")
|
||||
with open(cfg, "w", encoding="utf-8") as f:
|
||||
json.dump({"model": "m-from-cfg"}, f)
|
||||
self.assertEqual(_resolve_model(None, cfg), "m-from-cfg")
|
||||
with open(cfg, "w", encoding="utf-8") as f:
|
||||
f.write("{bad")
|
||||
self.assertEqual(_resolve_model(None, cfg), "gpt-5.6-terra"
|
||||
if _resolve_model(None, None) == "gpt-5.6-terra"
|
||||
else _resolve_model(None, None))
|
||||
self.assertEqual(_resolve_model(None, None),
|
||||
_resolve_model(None, os.path.join(
|
||||
tmp.name, "missing.json")))
|
||||
|
||||
def test_task_from_row(self):
|
||||
# idempotency_key/priority/last_heartbeat 是硬需列(load_tasks
|
||||
# 的显式列清单保证它们存在),其余新列走 _col 容错。
|
||||
row = {"id": "t1", "project_id": "p", "prompt": "x",
|
||||
"status": "queued", "depends_on": '["a", "b"]',
|
||||
"created_by": "o", "created_at": 1.0, "updated_at": 2.0,
|
||||
"idempotency_key": None, "priority": 0,
|
||||
"last_heartbeat": None}
|
||||
task = _task_from_row(row)
|
||||
self.assertEqual(task["depends_on"], ["a", "b"])
|
||||
row2 = dict(row, depends_on="not json")
|
||||
self.assertEqual(_task_from_row(row2)["depends_on"], [])
|
||||
row3 = dict(row, depends_on='"single"')
|
||||
self.assertEqual(_task_from_row(row3)["depends_on"], ["single"])
|
||||
row4 = dict(row, depends_on=None)
|
||||
self.assertEqual(_task_from_row(row4)["depends_on"], [])
|
||||
|
||||
|
||||
class WorkerDbTestCase(unittest.TestCase):
|
||||
def setUp(self):
|
||||
self.tmp = tempfile.TemporaryDirectory()
|
||||
self.addCleanup(self.tmp.cleanup)
|
||||
self.db_path = os.path.join(self.tmp.name, "control.db")
|
||||
init_db(self.db_path)
|
||||
self.hb = os.path.join(self.tmp.name, "hb.json")
|
||||
self.worker = Worker(self.db_path, "/fake/host.py",
|
||||
env={"XCM_API_KEY": "test-placeholder"},
|
||||
heartbeat_file=self.hb,
|
||||
worker_id="cov85-worker")
|
||||
|
||||
def _status(self, task_id):
|
||||
conn = sqlite3.connect(self.db_path)
|
||||
try:
|
||||
row = conn.execute("SELECT status FROM tasks WHERE id = ?",
|
||||
(task_id,)).fetchone()
|
||||
return row[0] if row else None
|
||||
finally:
|
||||
conn.close()
|
||||
|
||||
def test_code_signature_valid_and_changed(self):
|
||||
self.assertTrue(self.worker._code_signature_valid())
|
||||
self.assertFalse(self.worker._code_changed())
|
||||
# 启动签名被篡改后即判为已变且不可信(fail-closed 判据)。
|
||||
self.worker._code_signature = ("tampered",)
|
||||
self.assertTrue(self.worker._code_changed())
|
||||
self.assertFalse(self.worker._code_signature_valid())
|
||||
|
||||
def test_quota_hit_take_consumes(self):
|
||||
self.worker._mark_quota_hit("t1")
|
||||
self.assertTrue(self.worker._take_quota_hit("t1"))
|
||||
self.assertFalse(self.worker._take_quota_hit("t1")) # 取走即消费
|
||||
|
||||
def test_register_and_heartbeats(self):
|
||||
self.worker.register()
|
||||
self.worker._registry_heartbeat()
|
||||
self.worker._write_heartbeat()
|
||||
conn = sqlite3.connect(self.db_path)
|
||||
try:
|
||||
row = conn.execute("SELECT worker_id, last_seen FROM workers"
|
||||
" WHERE worker_id = ?",
|
||||
("cov85-worker",)).fetchone()
|
||||
finally:
|
||||
conn.close()
|
||||
self.assertIsNotNone(row)
|
||||
self.assertGreater(row[1], 0)
|
||||
with open(self.hb, encoding="utf-8") as f:
|
||||
payload = json.load(f)
|
||||
self.assertIn("last_seen", payload)
|
||||
self.assertEqual(payload["pid"], os.getpid())
|
||||
self.assertFalse(os.path.exists(self.hb + ".tmp")) # 原子替换无残留
|
||||
|
||||
def test_costs_lazy_and_close(self):
|
||||
store = self.worker.costs
|
||||
self.assertIsNotNone(store)
|
||||
self.assertIs(self.worker.costs, store) # 惰性单例
|
||||
self.worker.close() # 不抛错即可,注入的 store 由调用方管
|
||||
|
||||
def test_load_tasks_filters_and_get_task(self):
|
||||
_insert_task(self.db_path, "q1", status="queued")
|
||||
_insert_task(self.db_path, "r1", status="running")
|
||||
_insert_task(self.db_path, "c1", status="completed")
|
||||
_insert_task(self.db_path, "f1", status="failed")
|
||||
tasks = self.worker.load_tasks()
|
||||
self.assertEqual({t["id"] for t in tasks}, {"q1", "c1"})
|
||||
task = self.worker.get_task("q1")
|
||||
self.assertEqual(task["prompt"], "do it")
|
||||
self.assertIsNone(self.worker.get_task("missing"))
|
||||
|
||||
def test_claim_cas_and_capacity(self):
|
||||
_insert_task(self.db_path, "q1", status="queued")
|
||||
self.assertTrue(self.worker._claim("q1"))
|
||||
self.assertEqual(self._status("q1"), "running")
|
||||
self.assertFalse(self.worker._claim("q1")) # 已非 queued,CAS 失败
|
||||
self.worker._running_clear("q1")
|
||||
# 容量门:capacity=1 且已有占位时直接拒绝,不碰库。
|
||||
self.worker.capacity = 1
|
||||
self.worker._running_set("other")
|
||||
_insert_task(self.db_path, "q2", status="queued")
|
||||
self.assertFalse(self.worker._claim("q2"))
|
||||
self.assertEqual(self._status("q2"), "queued")
|
||||
self.worker._running_clear("other")
|
||||
|
||||
def test_claim_executor_filter(self):
|
||||
_insert_task(self.db_path, "d1", status="queued",
|
||||
executor="docker")
|
||||
_insert_task(self.db_path, "l1", status="queued",
|
||||
executor="local")
|
||||
_insert_task(self.db_path, "n1", status="queued",
|
||||
executor=None)
|
||||
# 本 worker 是 local:docker 任务不认领,local/NULL 认领。
|
||||
# 认领成功会占一个 _running 槽,容量默认 1,因此本测试
|
||||
# 先把容量放宽,否则第二单认领会被容量门挡住(非 CAS 问题)。
|
||||
self.worker.capacity = 10
|
||||
self.assertFalse(self.worker._claim("d1", executor="local"))
|
||||
self.assertTrue(self.worker._claim("l1", executor="local"))
|
||||
self.assertTrue(self.worker._claim("n1", executor="local"))
|
||||
# claim_executor_only 时 NULL 任务绝不认领。
|
||||
strict = Worker(self.db_path, "/fake/host.py",
|
||||
env={"XCM_API_KEY": "x"},
|
||||
claim_executor_only=True)
|
||||
_insert_task(self.db_path, "n2", status="queued", executor=None)
|
||||
self.assertFalse(strict._claim("n2", executor="local"))
|
||||
self.assertEqual(self._status("n2"), "queued")
|
||||
|
||||
def test_claim_refused_on_bad_signature(self):
|
||||
_insert_task(self.db_path, "q1", status="queued")
|
||||
self.worker._code_signature = ("tampered",)
|
||||
self.assertFalse(self.worker._claim("q1"))
|
||||
self.assertEqual(self._status("q1"), "queued")
|
||||
|
||||
def test_set_status_and_interrupted_reason(self):
|
||||
_insert_task(self.db_path, "t1", status="running")
|
||||
self.worker._set_status("t1", "completed")
|
||||
self.assertEqual(self._status("t1"), "completed")
|
||||
self.worker._mark_interrupted_reason("t1", "nope")
|
||||
conn = sqlite3.connect(self.db_path)
|
||||
try:
|
||||
row = conn.execute("SELECT interrupted_reason FROM tasks"
|
||||
" WHERE id = 't1'").fetchone()
|
||||
finally:
|
||||
conn.close()
|
||||
self.assertIsNone(row[0]) # 非 interrupted 任务不补原因
|
||||
_insert_task(self.db_path, "t2", status="interrupted")
|
||||
self.worker._mark_interrupted_reason("t2", "quota")
|
||||
conn = sqlite3.connect(self.db_path)
|
||||
try:
|
||||
row = conn.execute("SELECT interrupted_reason FROM tasks"
|
||||
" WHERE id = 't2'").fetchone()
|
||||
finally:
|
||||
conn.close()
|
||||
self.assertEqual(row[0], "quota")
|
||||
|
||||
def test_set_status_if_running_and_heartbeat(self):
|
||||
_insert_task(self.db_path, "t1", status="running")
|
||||
self.assertTrue(self.worker._set_status_if_running("t1", "failed"))
|
||||
self.assertFalse(self.worker._set_status_if_running("t1", "failed"))
|
||||
self.assertTrue(self.worker.heartbeat("t1") is False) # 已非 running
|
||||
_insert_task(self.db_path, "t2", status="running")
|
||||
self.assertTrue(self.worker.heartbeat("t2"))
|
||||
self.assertFalse(self.worker.heartbeat("missing"))
|
||||
|
||||
def test_set_child_pid(self):
|
||||
_insert_task(self.db_path, "t1", status="running")
|
||||
self.worker._set_child_pid("t1", 4242)
|
||||
conn = sqlite3.connect(self.db_path)
|
||||
try:
|
||||
row = conn.execute("SELECT child_pid FROM tasks"
|
||||
" WHERE id = 't1'").fetchone()
|
||||
finally:
|
||||
conn.close()
|
||||
self.assertEqual(row[0], 4242)
|
||||
|
||||
def test_kill_proc_group_real_process(self):
|
||||
proc = subprocess.Popen(["sleep", "60"], start_new_session=True)
|
||||
try:
|
||||
self.worker._kill_proc_group(proc)
|
||||
proc.wait(timeout=10)
|
||||
self.assertIsNotNone(proc.returncode)
|
||||
finally:
|
||||
if proc.poll() is None:
|
||||
proc.kill()
|
||||
# 已死进程再杀不抛错。
|
||||
self.worker._kill_proc_group(proc)
|
||||
|
||||
def test_running_bookkeeping(self):
|
||||
self.assertEqual(self.worker._running_count(), 0)
|
||||
self.worker._running_set("a", container="c1")
|
||||
self.worker._running_set("b")
|
||||
self.assertEqual(self.worker._running_count(), 2)
|
||||
snap = self.worker._running_snapshot()
|
||||
self.assertEqual(snap["a"]["container"], "c1")
|
||||
snap["a"]["container"] = "mutated"
|
||||
self.assertEqual(self.worker._running_snapshot()["a"]["container"],
|
||||
"c1") # 快照是深拷贝条目,不影响内部
|
||||
self.assertTrue(self.worker._wait_for_running(timeout=0.1) is False)
|
||||
self.worker._running_clear("a")
|
||||
self.worker._running_clear("b")
|
||||
self.assertTrue(self.worker._wait_for_running(timeout=1))
|
||||
# 心跳:running 任务刷 heartbeat,不抛错。
|
||||
_insert_task(self.db_path, "t1", status="running")
|
||||
self.worker._running_set("t1")
|
||||
self.worker._heartbeat_running_tasks()
|
||||
self.worker._running_clear("t1")
|
||||
|
||||
def test_cancel_paths(self):
|
||||
_insert_task(self.db_path, "q1", status="queued")
|
||||
self.assertTrue(self.worker.cancel("q1"))
|
||||
self.assertEqual(self._status("q1"), "interrupted")
|
||||
_insert_task(self.db_path, "done1", status="completed")
|
||||
self.assertFalse(self.worker.cancel("done1"))
|
||||
self.assertFalse(self.worker.cancel("missing"))
|
||||
# running + 真实子进程:cancel 先杀进程组再标状态。
|
||||
proc = subprocess.Popen(["sleep", "60"], start_new_session=True)
|
||||
try:
|
||||
_insert_task(self.db_path, "run1", status="running")
|
||||
self.worker._running_set("run1", proc=proc)
|
||||
self.assertTrue(self.worker.cancel("run1"))
|
||||
self.assertEqual(self._status("run1"), "interrupted")
|
||||
proc.wait(timeout=10)
|
||||
self.assertIsNotNone(proc.returncode)
|
||||
finally:
|
||||
if proc.poll() is None:
|
||||
proc.kill()
|
||||
self.worker._running_clear("run1")
|
||||
|
||||
def test_executor_name_key_and_popen(self):
|
||||
self.assertEqual(self.worker._executor_name(), "local")
|
||||
self.assertTrue(self.worker.key_available())
|
||||
no_key = Worker(self.db_path, "/fake/host.py", env={})
|
||||
self.assertFalse(no_key.key_available())
|
||||
self.assertTrue(self.worker._use_popen())
|
||||
factory_worker = Worker(self.db_path, "/fake/host.py",
|
||||
env={"XCM_API_KEY": "x"},
|
||||
process_factory=lambda *a, **k: None)
|
||||
self.assertTrue(factory_worker._use_popen())
|
||||
|
||||
def test_budget_and_quota(self):
|
||||
self.assertFalse(self.worker._budget_exceeded())
|
||||
self.worker.max_turns = 1
|
||||
self.worker.usage_turns = 1
|
||||
self.assertTrue(self.worker._budget_exceeded())
|
||||
self.worker.max_turns = None
|
||||
self.worker.usage_turns = 0
|
||||
# 无配额行 -> 默认配额;显式 NULL 行 -> 不限额。
|
||||
quota, used, month, unpriced = self.worker._quota_limit_and_used(
|
||||
"owner1")
|
||||
self.assertEqual(quota, 10000000)
|
||||
self.assertEqual(used, 0)
|
||||
conn = sqlite3.connect(self.db_path)
|
||||
conn.execute("INSERT INTO user_quotas(user_id, quota_usd_micros,"
|
||||
" updated_at) VALUES('owner1', NULL, 1.0)")
|
||||
conn.commit()
|
||||
conn.close()
|
||||
quota, used, month, unpriced = self.worker._quota_limit_and_used(
|
||||
"owner1")
|
||||
self.assertIsNone(quota)
|
||||
# 配额为 0 时 _quota_exceeded 返回超限文本;不限额返回 None。
|
||||
conn = sqlite3.connect(self.db_path)
|
||||
conn.execute("UPDATE user_quotas SET quota_usd_micros = 0"
|
||||
" WHERE user_id = 'owner1'")
|
||||
conn.commit()
|
||||
conn.close()
|
||||
task = {"id": "t1", "created_by": "owner1"}
|
||||
self.assertIsNotNone(self.worker._quota_exceeded(task, 0))
|
||||
conn = sqlite3.connect(self.db_path)
|
||||
conn.execute("UPDATE user_quotas SET quota_usd_micros = NULL"
|
||||
" WHERE user_id = 'owner1'")
|
||||
conn.commit()
|
||||
conn.close()
|
||||
self.assertIsNone(self.worker._quota_exceeded(task, 0))
|
||||
# 缺属主时按 fail-closed 文本(on_error 默认 open 时仍有文本?
|
||||
# 实测:缺属主直接返回 quota_exceeded 文本)。
|
||||
self.assertIsNotNone(self.worker._quota_exceeded({"id": "t9"}, 0))
|
||||
|
||||
def test_ready_and_pick_next(self):
|
||||
_insert_task(self.db_path, "a", status="queued")
|
||||
_insert_task(self.db_path, "b", status="queued",
|
||||
depends_on='["a"]')
|
||||
_insert_task(self.db_path, "c", status="queued",
|
||||
depends_on='["missing"]')
|
||||
ready = self.worker._ready_tasks()
|
||||
self.assertEqual([t["id"] for t in ready], ["a"])
|
||||
self.assertEqual(self.worker.pick_next()["id"], "a")
|
||||
# a 完成后 b 就绪。
|
||||
self.worker._set_status("a", "completed")
|
||||
ready = self.worker._ready_tasks()
|
||||
self.assertEqual([t["id"] for t in ready], ["b"])
|
||||
# 依赖成环 -> pick_next 为 None,不抛错。
|
||||
_insert_task(self.db_path, "x", status="queued",
|
||||
depends_on='["y"]')
|
||||
_insert_task(self.db_path, "y", status="queued",
|
||||
depends_on='["x"]')
|
||||
self.worker._set_status("b", "completed")
|
||||
self.assertIsNone(self.worker.pick_next())
|
||||
|
||||
def test_tenant_safe_name_and_env_dirs(self):
|
||||
self.assertEqual(self.worker._tenant_safe_name("a/b"), "a-b")
|
||||
with self.assertRaises(ValueError):
|
||||
self.worker._tenant_safe_name(None)
|
||||
with self.assertRaises(ValueError):
|
||||
self.worker._tenant_safe_name("...")
|
||||
self.worker.repos_dir = os.path.join(self.tmp.name, "repos")
|
||||
env = {"id": "env1", "project_id": "p1"}
|
||||
host_dir = self.worker._env_host_dir(env)
|
||||
self.assertTrue(host_dir.endswith(os.path.join("repos", "envs",
|
||||
"env1")))
|
||||
self.assertEqual(self.worker._env_repo_path(env), host_dir)
|
||||
|
||||
def test_fetch_environment_and_run_git(self):
|
||||
self.assertIsNone(self.worker._fetch_environment("missing"))
|
||||
conn = sqlite3.connect(self.db_path)
|
||||
try:
|
||||
from cloud.server import _ensure_environments_table
|
||||
_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)")
|
||||
conn.commit()
|
||||
finally:
|
||||
conn.close()
|
||||
env = self.worker._fetch_environment("e1")
|
||||
self.assertEqual(env["branch"], "main")
|
||||
self.assertEqual(env["status"], "ready")
|
||||
# 真实 git:空仓库里 rev-parse 应失败但返回结果对象。
|
||||
repo = os.path.join(self.tmp.name, "repo")
|
||||
os.makedirs(repo)
|
||||
result = self.worker._run_git(["git", "-C", repo, "rev-parse",
|
||||
"HEAD"])
|
||||
self.assertNotEqual(result.returncode, 0)
|
||||
# 注入的 repos_runner 两种签名都兼容。
|
||||
fake = Worker(self.db_path, "/fake/host.py",
|
||||
env={"XCM_API_KEY": "x"},
|
||||
repos_runner=lambda cmd: "ran")
|
||||
self.assertEqual(fake._run_git(["git", "status"]), "ran")
|
||||
|
||||
def two_arg(cmd, timeout=None):
|
||||
return ("two", timeout)
|
||||
|
||||
fake2 = Worker(self.db_path, "/fake/host.py",
|
||||
env={"XCM_API_KEY": "x"},
|
||||
repos_runner=lambda cmd, timeout=None: (_ for _ in ()).throw(TypeError("one-arg only")) if False else two_arg(cmd, timeout))
|
||||
# 上面 fake2 的 runner 本身接受两参,直接验证两参路径。
|
||||
self.assertEqual(fake2._run_git(["git"]), ("two", None))
|
||||
|
||||
def test_followup_claim_finish_restore_queue(self):
|
||||
_insert_task(self.db_path, "t1", status="running")
|
||||
conn = sqlite3.connect(self.db_path)
|
||||
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', 'hello', 'normal',"
|
||||
" 'pending', 'u1', 1.0, 1.0)")
|
||||
conn.commit()
|
||||
finally:
|
||||
conn.close()
|
||||
claimed = self.worker._claim_next_followup("t1")
|
||||
self.assertIsNotNone(claimed)
|
||||
self.assertEqual(claimed["id"], "f1")
|
||||
self.assertIsNone(self.worker._claim_next_followup("t1"))
|
||||
# restore 只把 running 放回 pending:先验证该路,再重认领
|
||||
# 并 finish 到 done。
|
||||
self.worker._restore_followup("f1")
|
||||
conn = sqlite3.connect(self.db_path)
|
||||
try:
|
||||
row = conn.execute("SELECT status FROM task_followups"
|
||||
" WHERE id = 'f1'").fetchone()
|
||||
finally:
|
||||
conn.close()
|
||||
self.assertEqual(row[0], "pending")
|
||||
self.assertIsNotNone(self.worker._claim_next_followup("t1"))
|
||||
self.worker._finish_followup("f1", "done")
|
||||
conn = sqlite3.connect(self.db_path)
|
||||
try:
|
||||
row = conn.execute("SELECT status FROM task_followups"
|
||||
" WHERE id = 'f1'").fetchone()
|
||||
finally:
|
||||
conn.close()
|
||||
self.assertEqual(row[0], "done")
|
||||
task = self.worker.get_task("t1")
|
||||
self.assertTrue(self.worker._queue_self_followup(task, "again"))
|
||||
conn = sqlite3.connect(self.db_path)
|
||||
try:
|
||||
rows = conn.execute("SELECT body FROM task_followups"
|
||||
" WHERE task_id = 't1' AND body = 'again'"
|
||||
).fetchall()
|
||||
finally:
|
||||
conn.close()
|
||||
self.assertEqual(len(rows), 1)
|
||||
|
||||
def test_never_started_and_breakpoint_resume(self):
|
||||
# 租户隔离 worker:state 路径在 tenants_root 下,全无痕迹
|
||||
# 时从未启动为 True;写 host-state 条目后为 False。
|
||||
tenant_worker = Worker(
|
||||
self.db_path, "/fake/host.py", env={"XCM_API_KEY": "x"},
|
||||
tenant_isolation=True,
|
||||
tenants_root=os.path.join(self.tmp.name, "tenants"))
|
||||
task = {"id": "t1", "project_id": "p1"}
|
||||
self.assertTrue(tenant_worker._task_never_started(task))
|
||||
state_dir = os.path.join(self.tmp.name, "tenants", "p1", "work")
|
||||
os.makedirs(state_dir, exist_ok=True)
|
||||
with open(os.path.join(state_dir, "host-state.json"),
|
||||
"w", encoding="utf-8") as f:
|
||||
json.dump({"tasks": {"t1": {}}}, f)
|
||||
self.assertFalse(tenant_worker._task_never_started(task))
|
||||
# 断点续跑:无断点 + 从未启动 -> 回落全新跑并清标记。
|
||||
fresh_task = {"id": "t9", "project_id": "p1",
|
||||
"_goal_resume_breakpoint": True}
|
||||
out = tenant_worker._resolve_breakpoint_resume(fresh_task)
|
||||
self.assertFalse(out["_goal_resume_breakpoint"])
|
||||
|
||||
def test_host_state_path_branches(self):
|
||||
# 租户路径。
|
||||
tenant_worker = Worker(
|
||||
self.db_path, "/fake/host.py", env={"XCM_API_KEY": "x"},
|
||||
tenant_isolation=True,
|
||||
tenants_root=os.path.join(self.tmp.name, "tenants"))
|
||||
self.assertEqual(
|
||||
tenant_worker._host_state_path_for_task(
|
||||
{"project_id": "p1"}),
|
||||
os.path.join(self.tmp.name, "tenants", "p1", "work",
|
||||
"host-state.json"))
|
||||
# 非租户无配置时回落默认路径。
|
||||
self.assertEqual(
|
||||
self.worker._host_state_path_for_task({"project_id": "p1"}),
|
||||
"/opt/codex-app-server/work/host-state.json")
|
||||
|
||||
def test_seed_dir_and_cleanup(self):
|
||||
task = {"id": "t1"}
|
||||
self.assertIsNone(self.worker._task_seed_dir(task))
|
||||
self.worker._cleanup_task_seed(task) # 无附件时静默通过
|
||||
bundle_dir = os.path.join(self.tmp.name, "uploads", "b1")
|
||||
os.makedirs(bundle_dir)
|
||||
with open(os.path.join(bundle_dir, "f.txt"), "w") as f:
|
||||
f.write("x")
|
||||
self.worker.upload_dir = os.path.join(self.tmp.name, "uploads")
|
||||
conn = sqlite3.connect(self.db_path)
|
||||
try:
|
||||
from cloud.server import _ensure_uploads_tables
|
||||
_ensure_uploads_tables(conn)
|
||||
conn.execute("INSERT INTO task_attachments(task_id, bundle_id,"
|
||||
" rel_path, size) VALUES('t1', 'b1', 'f.txt', 1)")
|
||||
conn.commit()
|
||||
finally:
|
||||
conn.close()
|
||||
self.assertEqual(self.worker._task_seed_dir(task), bundle_dir)
|
||||
self.worker._cleanup_task_seed(task)
|
||||
self.assertFalse(os.path.exists(bundle_dir))
|
||||
|
||||
def test_attempt_error_and_fallback_token(self):
|
||||
# 超时分支:result.status == "timeout" 时用统一超时文本。
|
||||
timed_out = Worker(self.db_path, "/fake/host.py",
|
||||
env={"XCM_API_KEY": "x"}, task_timeout=30)
|
||||
self.assertEqual(timed_out._attempt_error(
|
||||
{"status": "timeout"}, None, None),
|
||||
"task timeout after 30 seconds")
|
||||
self.assertEqual(self.worker._attempt_error({}, "plain err", None),
|
||||
"plain err")
|
||||
parsed = self.worker._attempt_error(
|
||||
{"stdout": json.dumps({"text": "from-result"})}, None, None)
|
||||
self.assertEqual(parsed, "from-result")
|
||||
# stdout JSON 里带 error 字段时取 error 文本。
|
||||
err_result = self.worker._attempt_error(
|
||||
{"stdout": json.dumps(
|
||||
{"error": {"message": "deep error"}})}, None, None)
|
||||
self.assertEqual(err_result, "deep error")
|
||||
# stdout 为空时取 stderr 尾部;两者皆空时退 token / 通用文本。
|
||||
self.assertEqual(self.worker._attempt_error(
|
||||
{"stdout": "", "stderr": " boom tail "}, None, None),
|
||||
"boom tail")
|
||||
self.assertEqual(self.worker._attempt_error({}, None, "404"),
|
||||
"host_status=404")
|
||||
self.assertEqual(self.worker._attempt_error({}, None, None),
|
||||
"host execution failed")
|
||||
self.assertEqual(self.worker._attempt_error(
|
||||
{"stdout": "plain stdout failure"}, None, None),
|
||||
"plain stdout failure")
|
||||
# _fallback_token:host_status 可 fallback 码直返;文本信号
|
||||
# 提取;无信号返回 None,绝不凭通用文本猜测。
|
||||
self.assertEqual(self.worker._fallback_token(
|
||||
{"host_status": "404"}, None), "404")
|
||||
self.assertEqual(self.worker._fallback_token(
|
||||
{"stderr": "model not found for this key"}, None), "404")
|
||||
self.assertEqual(self.worker._fallback_token(
|
||||
{"stdout": "upstream returned status 503"}, None), "503")
|
||||
self.assertIsNone(self.worker._fallback_token({}, None))
|
||||
self.assertIsNone(self.worker._fallback_token(
|
||||
{"host_status": "failed"}, None))
|
||||
|
||||
def test_coord_child_rows_and_orphans(self):
|
||||
_insert_task(self.db_path, "parent", status="running")
|
||||
_insert_task(self.db_path, "ch1", status="running")
|
||||
_insert_task(self.db_path, "ch2", status="completed")
|
||||
rows = self.worker._coord_child_rows(["ch1", "ch2", "gone"])
|
||||
self.assertEqual(rows["ch1"]["status"], "running")
|
||||
self.assertEqual(rows["ch2"]["status"], "completed")
|
||||
self.assertEqual(rows["gone"]["status"], "interrupted")
|
||||
marked = self.worker._coord_mark_orphans(
|
||||
{"id": "parent", "project_id": "proj"}, ["ch1", "ch2", "gone"])
|
||||
self.assertEqual(marked, ["ch1"])
|
||||
self.assertEqual(self._status("ch1"), "interrupted")
|
||||
self.assertEqual(self._status("ch2"), "completed")
|
||||
|
||||
def test_gateway_cache(self):
|
||||
# 默认网关(非 sub2api 环境)下 _gateway_context 返回空,
|
||||
# 且 _ensure_gateway 会把结果缓存进任务 dict。
|
||||
with _Env(CLOUD_GATEWAY=None):
|
||||
provider, env = self.worker._ensure_gateway(
|
||||
{"id": "t1", "created_by": "owner1"})
|
||||
self.assertIsNone(provider)
|
||||
self.assertIsNone(env)
|
||||
task = {"id": "t2", "_gateway_provider": "cached",
|
||||
"_gateway_env": {"k": "v"}}
|
||||
self.assertEqual(self.worker._ensure_gateway(task),
|
||||
("cached", {"k": "v"}))
|
||||
|
||||
def test_notify_task_finished_no_subscribers(self):
|
||||
# 无订阅时 dispatch 返回空 sent,不抛错、不影响任务。
|
||||
self.worker._notify_task_finished("t1", "proj", "completed")
|
||||
self.worker._notify_task_finished("t1", "proj", "failed",
|
||||
error="boom")
|
||||
|
||||
def test_audit_best_effort(self):
|
||||
self.worker._audit("proj", "cov85_probe", "detail")
|
||||
conn = sqlite3.connect(self.db_path)
|
||||
try:
|
||||
row = conn.execute("SELECT COUNT(*) FROM audit_log"
|
||||
" WHERE action = 'cov85_probe'").fetchone()
|
||||
finally:
|
||||
conn.close()
|
||||
self.assertEqual(row[0], 1)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
File diff suppressed because it is too large.
Load diff
Reference in new issue
Block a user