host
488aedbcad
(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).
706 lines
31 KiB
Python
706 lines
31 KiB
Python
"""覆盖率 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 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()
|