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).
540 lines
21 KiB
Python
540 lines
21 KiB
Python
"""Coverage batch v4-B1: cloud/worker.py small pure/error branches.
|
||
|
||
Stubbing discipline: the method under test is never stubbed. Patched
|
||
collaborators are: register_worker / heartbeat_worker (registry sinks),
|
||
AuthStore (audit sink), AutoMemoryStore / _auto_prefix / _todo_list_db /
|
||
_db_connect (data collaborators), log_event (only to drive the
|
||
log-failure swallow branches), os.killpg (to drive the kill-error
|
||
branches without depending on a real process group dying mid-test),
|
||
and instance-level _host_state_path_for_task / _rebuild... /
|
||
_task_never_started / _fetch_environment predicates where the real
|
||
path cannot produce the branch without a race or a missing row.
|
||
Every assertion reads a real return value or a real DB row.
|
||
"""
|
||
|
||
import os
|
||
import signal
|
||
import sqlite3
|
||
import subprocess
|
||
import sys
|
||
import tempfile
|
||
import time
|
||
import unittest
|
||
from pathlib import Path
|
||
from types import SimpleNamespace
|
||
from unittest import mock
|
||
|
||
REPO_ROOT = Path(__file__).resolve().parents[2]
|
||
sys.path.insert(0, str(REPO_ROOT))
|
||
|
||
from cloud.server import init_db # noqa: E402
|
||
from cloud.worker import Worker, _code_signature # noqa: E402
|
||
|
||
FAKE_ENV = {"XCM_API_KEY": "test-placeholder"}
|
||
|
||
|
||
def _worker(db_path, **kwargs):
|
||
return Worker(db_path, "/fake/host.py", env=FAKE_ENV, **kwargs)
|
||
|
||
|
||
class _Base(unittest.TestCase):
|
||
def setUp(self):
|
||
self.tmp = tempfile.TemporaryDirectory()
|
||
self.db_path = os.path.join(self.tmp.name, "control.db")
|
||
init_db(self.db_path)
|
||
|
||
def tearDown(self):
|
||
self.tmp.cleanup()
|
||
|
||
def _insert_task(self, task_id="t1", status="queued", **cols):
|
||
conn = sqlite3.connect(self.db_path)
|
||
try:
|
||
conn.execute(
|
||
"INSERT INTO tasks(id, project_id, prompt, status,"
|
||
" created_by, created_at, updated_at)"
|
||
" VALUES(?, ?, ?, ?, ?, ?, ?)",
|
||
(task_id, "p1", "do work", status, "u1", time.time(),
|
||
time.time()))
|
||
for key, value in cols.items():
|
||
conn.execute(
|
||
"UPDATE tasks SET %s = ? WHERE id = ?" % key,
|
||
(value, task_id))
|
||
conn.commit()
|
||
finally:
|
||
conn.close()
|
||
|
||
def _row(self, task_id="t1"):
|
||
conn = sqlite3.connect(self.db_path)
|
||
try:
|
||
conn.row_factory = sqlite3.Row
|
||
return conn.execute(
|
||
"SELECT * FROM tasks WHERE id = ?", (task_id,)).fetchone()
|
||
finally:
|
||
conn.close()
|
||
|
||
|
||
class InitAndRegistryTestCase(_Base):
|
||
def test_code_signature_stat_failure_none(self):
|
||
with mock.patch.object(os, "stat", side_effect=OSError("gone")):
|
||
self.assertIsNone(_code_signature("/nonexistent/host.py"))
|
||
|
||
def test_init_router_type_error(self):
|
||
with self.assertRaises(ValueError):
|
||
_worker(self.db_path, router="not-a-router")
|
||
|
||
def test_init_fallback_models_normalized(self):
|
||
w = _worker(self.db_path, model="m0",
|
||
fallback_models=[123, " m1 ", "m1", "m0", "m2"])
|
||
self.assertEqual(w.fallback_models, ["m1", "m2"])
|
||
|
||
def test_init_executor_budget_flag_failure_swallowed(self):
|
||
class _BadExecutor:
|
||
@property
|
||
def ignore_host_budget(self):
|
||
return False
|
||
|
||
@ignore_host_budget.setter
|
||
def ignore_host_budget(self, value):
|
||
raise RuntimeError("read-only")
|
||
w = _worker(self.db_path, executor=_BadExecutor())
|
||
self.assertIsInstance(w.executor, _BadExecutor)
|
||
|
||
def test_register_failure_swallowed(self):
|
||
w = _worker(self.db_path)
|
||
with mock.patch("cloud.worker.register_worker",
|
||
side_effect=RuntimeError("db down")):
|
||
w.register()
|
||
|
||
def test_registry_heartbeat_failure_swallowed(self):
|
||
w = _worker(self.db_path)
|
||
with mock.patch("cloud.worker.heartbeat_worker",
|
||
side_effect=RuntimeError("db down")):
|
||
w._registry_heartbeat()
|
||
|
||
def test_close_cost_store_failure_swallowed(self):
|
||
w = _worker(self.db_path)
|
||
store = mock.Mock()
|
||
store._conn.close.side_effect = RuntimeError("down")
|
||
w._cost_store = store
|
||
w.close()
|
||
|
||
|
||
class StepAndCheckpointTestCase(_Base):
|
||
def test_mark_interrupted_reason_db_failure_swallowed(self):
|
||
w = _worker(self.tmp.name)
|
||
w._mark_interrupted_reason("t1", "quota")
|
||
|
||
def test_record_step_conn_failure_swallowed(self):
|
||
w = _worker(self.db_path)
|
||
conn = mock.Mock()
|
||
conn.execute.side_effect = RuntimeError("db down")
|
||
w._record_step(conn, "t1", "round", "hello")
|
||
|
||
def test_checkpoint_round_real_and_failure(self):
|
||
w = _worker(self.db_path)
|
||
conn = sqlite3.connect(self.db_path)
|
||
try:
|
||
conn.execute(
|
||
"CREATE TABLE IF NOT EXISTS task_checkpoints("
|
||
"task_id TEXT, round_no INTEGER, thread_id TEXT,"
|
||
" goal_status TEXT, state_sha TEXT, verify_out TEXT,"
|
||
" created_at DOUBLE PRECISION,"
|
||
" PRIMARY KEY(task_id, round_no))")
|
||
conn.execute("INSERT INTO task_checkpoints VALUES"
|
||
"('t1', 2, 'th', 'ok', 'sha', '', 0)")
|
||
conn.execute("INSERT INTO task_checkpoints VALUES"
|
||
"('t1', 5, 'th', 'ok', 'sha', '', 0)")
|
||
conn.commit()
|
||
self.assertEqual(w._checkpoint_round(conn, "t1"), 5)
|
||
self.assertEqual(w._checkpoint_round(conn, "nobody"), -1)
|
||
finally:
|
||
conn.close()
|
||
bad = mock.Mock()
|
||
bad.execute.side_effect = RuntimeError("db down")
|
||
self.assertEqual(w._checkpoint_round(bad, "t1"), -1)
|
||
|
||
def test_checkpoint_mirror_loop_exception_swallowed(self):
|
||
w = _worker(self.db_path)
|
||
w._mirror_checkpoint_once = lambda task: (_ for _ in ()).throw(
|
||
RuntimeError("mirror down"))
|
||
class _Stop:
|
||
def __init__(self):
|
||
self.n = 0
|
||
|
||
def wait(self, timeout):
|
||
self.n += 1
|
||
return self.n >= 2
|
||
w._checkpoint_mirror_loop({"id": "t1"}, _Stop())
|
||
|
||
def test_host_state_path_mounts_tuple(self):
|
||
w = _worker(self.db_path)
|
||
w.executor = SimpleNamespace(extra_mounts=[("/srv/work",
|
||
"/app/work")],
|
||
config_path=None)
|
||
self.assertEqual(
|
||
w._host_state_path_for_task({"project_id": "p1", "id": "t1"}),
|
||
"/srv/work/host-state.json")
|
||
|
||
def test_host_state_path_non_tuple_mount_and_none_task(self):
|
||
w = _worker(self.db_path)
|
||
w.executor = SimpleNamespace(extra_mounts=[123], config_path=None)
|
||
self.assertEqual(
|
||
w._host_state_path_for_task({"project_id": "p1", "id": "t1"}),
|
||
"/opt/codex-app-server/work/host-state.json")
|
||
self.assertIsNone(w._host_state_path_for_task(None))
|
||
|
||
def test_mirror_checkpoint_once_bad_json(self):
|
||
w = _worker(self.db_path)
|
||
work = os.path.join(self.tmp.name, "work")
|
||
os.makedirs(work, exist_ok=True)
|
||
with open(os.path.join(work, "host-state.json"), "w") as fh:
|
||
fh.write("{not json")
|
||
w.executor = SimpleNamespace(extra_mounts=[(work, "/app/work")],
|
||
config_path=None)
|
||
self.assertEqual(
|
||
w._mirror_checkpoint_once({"project_id": "p1", "id": "t1"}), -1)
|
||
|
||
def test_import_host_steps_task_none_swallowed(self):
|
||
w = _worker(self.db_path)
|
||
conn = sqlite3.connect(self.db_path)
|
||
try:
|
||
w._import_host_steps(conn, None, "/tmp/x/host-state.json")
|
||
finally:
|
||
conn.close()
|
||
|
||
def test_rebuild_trow_missing(self):
|
||
w = _worker(self.db_path)
|
||
conn = sqlite3.connect(self.db_path)
|
||
try:
|
||
conn.execute(
|
||
"CREATE TABLE IF NOT EXISTS task_checkpoints("
|
||
"task_id TEXT, round_no INTEGER, thread_id TEXT,"
|
||
" goal_status TEXT, state_sha TEXT, verify_out TEXT,"
|
||
" created_at DOUBLE PRECISION,"
|
||
" PRIMARY KEY(task_id, round_no))")
|
||
conn.execute("INSERT INTO task_checkpoints VALUES"
|
||
"('ghost', 3, 'th', 'ok', 'sha', '', 0)")
|
||
conn.commit()
|
||
finally:
|
||
conn.close()
|
||
self.assertFalse(w._rebuild_host_state_from_checkpoint(
|
||
{"id": "ghost", "project_id": "p1"}))
|
||
|
||
def test_rebuild_state_path_none(self):
|
||
self._insert_task()
|
||
w = _worker(self.db_path)
|
||
conn = sqlite3.connect(self.db_path)
|
||
try:
|
||
conn.execute(
|
||
"CREATE TABLE IF NOT EXISTS task_checkpoints("
|
||
"task_id TEXT, round_no INTEGER, thread_id TEXT,"
|
||
" goal_status TEXT, state_sha TEXT, verify_out TEXT,"
|
||
" created_at DOUBLE PRECISION,"
|
||
" PRIMARY KEY(task_id, round_no))")
|
||
conn.execute("INSERT INTO task_checkpoints VALUES"
|
||
"('t1', 3, 'th', 'ok', 'sha', '', 0)")
|
||
conn.commit()
|
||
finally:
|
||
conn.close()
|
||
w._host_state_path_for_task = lambda task: None
|
||
self.assertFalse(w._rebuild_host_state_from_checkpoint(
|
||
{"id": "t1", "project_id": "p1"}))
|
||
|
||
def test_rebuild_already_present_true(self):
|
||
self._insert_task()
|
||
w = _worker(self.db_path)
|
||
conn = sqlite3.connect(self.db_path)
|
||
try:
|
||
conn.execute(
|
||
"CREATE TABLE IF NOT EXISTS task_checkpoints("
|
||
"task_id TEXT, round_no INTEGER, thread_id TEXT,"
|
||
" goal_status TEXT, state_sha TEXT, verify_out TEXT,"
|
||
" created_at DOUBLE PRECISION,"
|
||
" PRIMARY KEY(task_id, round_no))")
|
||
conn.execute("INSERT INTO task_checkpoints VALUES"
|
||
"('t1', 3, 'th', 'ok', 'sha', '', 0)")
|
||
conn.commit()
|
||
finally:
|
||
conn.close()
|
||
work = os.path.join(self.tmp.name, "work")
|
||
os.makedirs(work, exist_ok=True)
|
||
with open(os.path.join(work, "host-state.json"), "w") as fh:
|
||
import json as _json
|
||
_json.dump({"version": 1, "tasks": {"t1": {"goal": "g"}}}, fh)
|
||
w.executor = SimpleNamespace(extra_mounts=[(work, "/app/work")],
|
||
config_path=None)
|
||
self.assertTrue(w._rebuild_host_state_from_checkpoint(
|
||
{"id": "t1", "project_id": "p1"}))
|
||
|
||
def test_rebuild_db_failure_false(self):
|
||
w = _worker(self.tmp.name)
|
||
self.assertFalse(w._rebuild_host_state_from_checkpoint(
|
||
{"id": "t1", "project_id": "p1"}))
|
||
|
||
def test_never_started_state_path_none(self):
|
||
w = _worker(self.db_path)
|
||
w._host_state_path_for_task = lambda task: None
|
||
self.assertFalse(w._task_never_started(
|
||
{"id": "t1", "project_id": "p1"}))
|
||
|
||
def test_resolve_breakpoint_log_failure_swallowed(self):
|
||
w = _worker(self.db_path)
|
||
w._rebuild_host_state_from_checkpoint = lambda task: (_ for _ in (
|
||
)).throw(RuntimeError("rebuild down"))
|
||
w._task_never_started = lambda task: True
|
||
with mock.patch("cloud.worker.log_event",
|
||
side_effect=RuntimeError("log down")):
|
||
out = w._resolve_breakpoint_resume(
|
||
{"id": "t1", "project_id": "p1"})
|
||
self.assertFalse(out["_goal_resume_breakpoint"])
|
||
|
||
|
||
class EnsureColumnsTestCase(_Base):
|
||
def _fake(self, mode):
|
||
all_cols = ["result_text", "result_text_full", "error_text",
|
||
"actual_model", "started_at", "finished_at",
|
||
"child_pid", "executor", "pr_url", "pr_number", "goal",
|
||
"verify_command", "goal_max_rounds", "goal_mode",
|
||
"worktree_host", "claimed_by_worker", "model", "effort",
|
||
"child_model", "child_effort", "approval", "plan",
|
||
"forked_from", "environment_id", "resume_message",
|
||
"turn_id", "root_turn_id", "coordinator_parts",
|
||
"parent_task_id", "coordinator_part",
|
||
"goal_resume_pending", "paused_by_user",
|
||
"interrupted_reason", "auto_resume_count",
|
||
"no_time_limit", "resume_after",
|
||
"auto_resume_stall_count",
|
||
"checkpoint_round_at_last_resume",
|
||
"last_error_signature"]
|
||
|
||
class _Fake:
|
||
def execute(self, sql, params=()):
|
||
if sql.strip().startswith("PRAGMA"):
|
||
rows = [(i, name, "TEXT", 0, None, 0)
|
||
for i, name in enumerate(
|
||
all_cols if mode == "create" else ["id"])]
|
||
result = mock.Mock()
|
||
result.fetchall.return_value = rows
|
||
return result
|
||
if sql.strip().startswith("ALTER"):
|
||
raise RuntimeError("alter down")
|
||
if sql.strip().startswith("CREATE"):
|
||
if mode == "create":
|
||
raise RuntimeError("create down")
|
||
return mock.Mock()
|
||
return mock.Mock()
|
||
return _Fake()
|
||
|
||
def test_ensure_columns_pragma_failure(self):
|
||
w = _worker(self.db_path)
|
||
bad = mock.Mock()
|
||
bad.execute.side_effect = RuntimeError("pragma down")
|
||
w._ensure_result_columns(bad)
|
||
|
||
def test_ensure_columns_alter_failure(self):
|
||
w = _worker(self.db_path)
|
||
w._ensure_result_columns(self._fake("alter"))
|
||
|
||
def test_ensure_columns_create_failure(self):
|
||
w = _worker(self.db_path)
|
||
w._ensure_result_columns(self._fake("create"))
|
||
|
||
|
||
class WriteResultTestCase(_Base):
|
||
def test_write_result_child_pid_branch(self):
|
||
self._insert_task()
|
||
w = _worker(self.db_path)
|
||
self.assertTrue(w._write_execution_result(
|
||
"t1", result_text="done", child_pid=99))
|
||
self.assertEqual(self._row()["child_pid"], 99)
|
||
|
||
def test_write_result_old_column_fallback(self):
|
||
self._insert_task()
|
||
w = _worker(self.db_path)
|
||
w._ensure_result_columns = lambda conn: None
|
||
real = sqlite3.connect(self.db_path)
|
||
|
||
class _Flaky:
|
||
def __init__(self):
|
||
self.flipped = False
|
||
|
||
def execute(self, sql, params=()):
|
||
if "result_text_full" in sql and not self.flipped:
|
||
self.flipped = True
|
||
raise sqlite3.OperationalError(
|
||
"no such column: result_text_full")
|
||
return real.execute(sql, params)
|
||
|
||
def rollback(self):
|
||
raise RuntimeError("rollback down")
|
||
|
||
def commit(self):
|
||
return real.commit()
|
||
|
||
def close(self):
|
||
return real.close()
|
||
with mock.patch("cloud.worker._db_connect",
|
||
return_value=_Flaky()):
|
||
self.assertTrue(w._write_execution_result(
|
||
"t1", result_text="fallback-ok"))
|
||
self.assertEqual(self._row()["result_text"], "fallback-ok")
|
||
|
||
|
||
class KillAuditNotifyTestCase(_Base):
|
||
def test_kill_proc_group_term_error(self):
|
||
w = _worker(self.db_path)
|
||
proc = subprocess.Popen(["sleep", "30"], start_new_session=True)
|
||
try:
|
||
with mock.patch("os.killpg",
|
||
side_effect=ProcessLookupError("gone")):
|
||
w._kill_proc_group(proc)
|
||
self.assertIsNone(proc.poll())
|
||
finally:
|
||
proc.kill()
|
||
proc.wait()
|
||
|
||
def test_kill_proc_group_sigkill_error(self):
|
||
w = _worker(self.db_path)
|
||
proc = subprocess.Popen(["sleep", "30"], start_new_session=True)
|
||
try:
|
||
def _killpg(pgid, sig):
|
||
if sig == signal.SIGKILL:
|
||
raise ProcessLookupError("gone")
|
||
return None
|
||
with mock.patch("os.killpg", side_effect=_killpg):
|
||
w._kill_proc_group(proc)
|
||
self.assertIsNone(proc.poll())
|
||
finally:
|
||
proc.kill()
|
||
proc.wait()
|
||
|
||
def test_audit_close_failure_swallowed(self):
|
||
w = _worker(self.db_path)
|
||
|
||
class _Store:
|
||
def __init__(self, db_path):
|
||
self.conn = mock.Mock()
|
||
self.conn.close.side_effect = RuntimeError("down")
|
||
|
||
def audit(self, *args):
|
||
return None
|
||
with mock.patch("cloud.worker.AuthStore", _Store):
|
||
w._audit("p1", "act")
|
||
|
||
def test_notify_errors_and_exception_swallowed(self):
|
||
w = _worker(self.db_path)
|
||
with mock.patch("cloud.push.dispatch_event",
|
||
return_value={"errors": ["e1"]}):
|
||
w._notify_task_finished("t1", "p1", "completed")
|
||
with mock.patch("cloud.push.dispatch_event",
|
||
side_effect=RuntimeError("push down")):
|
||
w._notify_task_finished("t1", "p1", "failed")
|
||
|
||
|
||
class InjectionTestCase(_Base):
|
||
def test_auto_memory_no_owner(self):
|
||
w = _worker(self.db_path)
|
||
task = {"id": "t1", "project_id": "p1", "prompt": "p"}
|
||
self.assertIs(w._apply_auto_memory_injection(task), task)
|
||
|
||
def test_auto_memory_store_and_log_failure(self):
|
||
w = _worker(self.db_path)
|
||
w.logger = mock.Mock()
|
||
w.logger.warning.side_effect = RuntimeError("log down")
|
||
task = {"id": "t1", "project_id": "p1", "prompt": "p",
|
||
"created_by": "u1"}
|
||
with mock.patch("cloud.worker.AutoMemoryStore",
|
||
side_effect=RuntimeError("store down")):
|
||
self.assertIs(w._apply_auto_memory_injection(task), task)
|
||
|
||
def test_auto_memory_goal_branch(self):
|
||
w = _worker(self.db_path)
|
||
|
||
class _Store:
|
||
def __init__(self, db_path):
|
||
pass
|
||
|
||
def get_task_flags(self, task_id):
|
||
return {}
|
||
|
||
def close(self):
|
||
pass
|
||
task = {"id": "t1", "project_id": "p1", "prompt": "P",
|
||
"goal": "G", "goal_mode": 1, "created_by": "u1"}
|
||
with mock.patch("cloud.worker.AutoMemoryStore", _Store), \
|
||
mock.patch("cloud.worker._auto_prefix",
|
||
return_value="MEM"):
|
||
out = w._apply_auto_memory_injection(task)
|
||
self.assertEqual(out["goal"], "MEM\n\nG")
|
||
self.assertEqual(out["prompt"], "MEM\n\nP")
|
||
self.assertTrue(out["_memory_injected"])
|
||
self.assertEqual(task["goal"], "G")
|
||
|
||
def test_todo_injection_db_and_log_failure(self):
|
||
w = _worker(self.tmp.name)
|
||
w.logger = mock.Mock()
|
||
w.logger.warning.side_effect = RuntimeError("log down")
|
||
task = {"id": "t1", "project_id": "p1", "prompt": "p"}
|
||
self.assertIs(w._apply_todo_injection(task), task)
|
||
|
||
def test_todo_injection_over_twenty(self):
|
||
w = _worker(self.db_path)
|
||
todos = [{"id": "td%d" % i, "title": "T%d" % i,
|
||
"status": "pending"} for i in range(21)]
|
||
task = {"id": "t1", "project_id": "p1", "prompt": "p"}
|
||
with mock.patch("cloud.worker._todo_list_db",
|
||
return_value=todos):
|
||
out = w._apply_todo_injection(task)
|
||
self.assertTrue(out["_todo_injected"])
|
||
self.assertIn("还有 1 条未列出", out["prompt"])
|
||
|
||
def test_project_injection_db_and_log_failure(self):
|
||
w = _worker(self.tmp.name)
|
||
w.logger = mock.Mock()
|
||
w.logger.warning.side_effect = RuntimeError("log down")
|
||
task = {"id": "t1", "project_id": "p1", "prompt": "p"}
|
||
self.assertIs(w._apply_project_injection(task), task)
|
||
|
||
def test_project_injection_row_key_errors(self):
|
||
w = _worker(self.db_path)
|
||
|
||
class _BadRow:
|
||
def __getitem__(self, key):
|
||
raise KeyError(key)
|
||
|
||
class _FakeConn:
|
||
def execute(self, sql, params=()):
|
||
result = mock.Mock()
|
||
result.fetchone.return_value = _BadRow()
|
||
return result
|
||
|
||
def close(self):
|
||
pass
|
||
task = {"id": "t1", "project_id": "p1", "prompt": "p"}
|
||
with mock.patch("cloud.worker._db_connect",
|
||
return_value=_FakeConn()):
|
||
out = w._apply_project_injection(task)
|
||
self.assertTrue(out["_project_injected"])
|
||
self.assertIn("项目：p1（id=p1）", out["prompt"])
|
||
self.assertIn("仓库：未绑定", out["prompt"])
|
||
|
||
def test_project_injection_env_missing(self):
|
||
conn = sqlite3.connect(self.db_path)
|
||
try:
|
||
conn.execute(
|
||
"INSERT INTO projects(id, name, owner_id, created_at,"
|
||
" repo_url, auto_pr) VALUES('p1', 'Proj', 'u1', 0, NULL, 0)")
|
||
conn.commit()
|
||
finally:
|
||
conn.close()
|
||
w = _worker(self.db_path)
|
||
w._fetch_environment = lambda env_id: None
|
||
task = {"id": "t1", "project_id": "p1", "prompt": "p",
|
||
"environment_id": "envX"}
|
||
out = w._apply_project_injection(task)
|
||
self.assertIn("envX（查不到该环境）", out["prompt"])
|
||
|
||
|
||
if __name__ == "__main__":
|
||
unittest.main()
|