(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).
302 lines
11 KiB
Python
302 lines
11 KiB
Python
"""Coverage batch v4-D1: cloud/worker.py resume/checkpoint/notify branches.
|
|
|
|
Stubbing discipline: the method under test is never stubbed. Patched
|
|
collaborators are: _quota_exceeded / _attempt_mark / _db_connect /
|
|
dispatch_event / AutoMemoryStore / the instance logger (data and sink
|
|
collaborators), plus instance-level _mirror_checkpoint_once where the
|
|
real path cannot raise without a race. Every assertion reads a real
|
|
return value or a real DB row.
|
|
"""
|
|
|
|
import json
|
|
import os
|
|
import sqlite3
|
|
import sys
|
|
import tempfile
|
|
import threading
|
|
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 # 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()
|
|
|
|
def _fake_executor(self, state_dir):
|
|
cfg = os.path.join(state_dir, "config.json")
|
|
with open(cfg, "w", encoding="utf-8") as fh:
|
|
json.dump({"state": "host-state.json"}, fh)
|
|
return SimpleNamespace(extra_mounts=[], config_path=cfg)
|
|
|
|
|
|
class AutoResumeExceptionTestCase(_Base):
|
|
def test_quota_check_raises_stays_interrupted(self):
|
|
# 1093-1094: _quota_exceeded raising inside auto_resume_interrupted
|
|
# must be swallowed into "quota_check_failed" and the task skipped.
|
|
self._insert_task(status="interrupted", goal_mode=1,
|
|
interrupted_reason="quota")
|
|
w = _worker(self.db_path)
|
|
with mock.patch.object(
|
|
Worker, "_quota_exceeded",
|
|
side_effect=RuntimeError("quota db down")):
|
|
resumed = w.auto_resume_interrupted()
|
|
self.assertNotIn("t1", resumed)
|
|
self.assertEqual(self._row()["status"], "interrupted")
|
|
|
|
def test_attempt_mark_failure_swallowed_on_resume(self):
|
|
# 1163-1164: _attempt_mark raising after a resume must not block
|
|
# the resume itself.
|
|
self._insert_task(status="interrupted", goal_mode=1,
|
|
interrupted_reason="worker_shutdown")
|
|
w = _worker(self.db_path, attempt_ledger_hook=True)
|
|
with mock.patch("cloud.worker._attempt_mark",
|
|
side_effect=RuntimeError("ledger down")):
|
|
resumed = w.auto_resume_interrupted()
|
|
self.assertIn("t1", resumed)
|
|
self.assertEqual(self._row()["status"], "queued")
|
|
|
|
def test_db_connect_failure_returns_empty(self):
|
|
# 1182-1183: outer except of auto_resume_interrupted.
|
|
w = _worker(self.db_path)
|
|
with mock.patch("cloud.worker._db_connect",
|
|
side_effect=RuntimeError("db down")):
|
|
self.assertEqual(w.auto_resume_interrupted(), [])
|
|
|
|
def test_checkpoint_select_failure_returns_minus_one(self):
|
|
# 1102-1103: checkpoint SELECT raising inside auto_resume must
|
|
# fall back to cp_round=-1 and still resume a platform task.
|
|
# NOTE: dropping the table cannot reach this branch, because
|
|
# _ensure_result_columns recreates task_checkpoints before the
|
|
# SELECT; the branch is reached only when the SELECT itself
|
|
# raises, so the connection is faked for that statement only.
|
|
self._insert_task(status="interrupted", goal_mode=1,
|
|
interrupted_reason="worker_shutdown")
|
|
w = _worker(self.db_path)
|
|
w._ensure_result_columns = lambda conn: None
|
|
row = ("t1", "p1", "u1", None, "m0", "worker_shutdown", None,
|
|
0, None, 1, None, None)
|
|
|
|
class _Cursor:
|
|
def __init__(self, rows=None, rowcount=1):
|
|
self._rows = rows or []
|
|
self.rowcount = rowcount
|
|
|
|
def fetchall(self):
|
|
return self._rows
|
|
|
|
def fetchone(self):
|
|
return self._rows[0] if self._rows else None
|
|
|
|
class _FakeConn:
|
|
def execute(self, sql, params=()):
|
|
if "task_checkpoints" in sql:
|
|
raise RuntimeError("checkpoint read down")
|
|
if "FROM tasks" in sql and "SELECT" in sql:
|
|
return _Cursor(rows=[row])
|
|
return _Cursor()
|
|
|
|
def commit(self):
|
|
pass
|
|
|
|
def close(self):
|
|
pass
|
|
|
|
with mock.patch("cloud.worker._db_connect",
|
|
return_value=_FakeConn()):
|
|
resumed = w.auto_resume_interrupted()
|
|
self.assertIn("t1", resumed)
|
|
|
|
def test_checkpoint_round_and_mirror_loop_exceptions(self):
|
|
w = _worker(self.db_path)
|
|
conn = sqlite3.connect(self.db_path)
|
|
try:
|
|
self.assertEqual(w._checkpoint_round(conn, "nope"), -1)
|
|
finally:
|
|
conn.close()
|
|
bad = mock.Mock()
|
|
bad.execute.side_effect = RuntimeError("boom")
|
|
self.assertEqual(w._checkpoint_round(bad, "nope"), -1)
|
|
# 1230-1231: first mirror call raising, stop already set.
|
|
w._mirror_checkpoint_once = mock.Mock(
|
|
side_effect=RuntimeError("mirror down"))
|
|
stop = threading.Event()
|
|
stop.set()
|
|
w._checkpoint_mirror_loop({"id": "t1"}, stop)
|
|
# 1233-1236: loop-body mirror call raising, then stop.
|
|
calls = {"n": 0}
|
|
|
|
def _wait(timeout=None):
|
|
calls["n"] += 1
|
|
return calls["n"] > 1
|
|
|
|
class _Stop:
|
|
def wait(self, timeout=None):
|
|
return _wait(timeout)
|
|
|
|
w._checkpoint_mirror_loop({"id": "t1"}, _Stop())
|
|
|
|
|
|
class HostStatePathTestCase(_Base):
|
|
def test_non_tuple_mount_item_falls_through(self):
|
|
# 1256: a non-tuple mounts entry yields (None, None) and the
|
|
# loop falls through to the /opt fallback.
|
|
w = _worker(self.db_path, executor=SimpleNamespace(
|
|
extra_mounts=["not-a-tuple"], config_path=None))
|
|
path = w._host_state_path_for_task({"project_id": "p1"})
|
|
self.assertEqual(path, "/opt/codex-app-server/work/host-state.json")
|
|
|
|
def test_mirror_missing_state_returns_minus_one(self):
|
|
# 1286: no state file at all.
|
|
state_dir = os.path.join(self.tmp.name, "cfg")
|
|
os.makedirs(state_dir)
|
|
w = _worker(self.db_path,
|
|
executor=self._fake_executor(state_dir))
|
|
self.assertEqual(
|
|
w._mirror_checkpoint_once({"id": "t1", "project_id": "p1"}), -1)
|
|
|
|
def test_rebuild_from_checkpoint_missing_state_file(self):
|
|
# 1411: state file absent -> fresh {"version": 1, "tasks": {}}
|
|
# is built and the rebuilt entry is written back.
|
|
self._insert_task(status="interrupted", verify_command="true",
|
|
goal_max_rounds=5)
|
|
conn = sqlite3.connect(self.db_path)
|
|
try:
|
|
conn.execute(
|
|
"INSERT INTO task_checkpoints(task_id, round_no,"
|
|
" thread_id, goal_status) VALUES(?, ?, ?, ?)",
|
|
("t1", 2, "th1", "interrupted"))
|
|
conn.commit()
|
|
finally:
|
|
conn.close()
|
|
state_dir = os.path.join(self.tmp.name, "cfg2")
|
|
os.makedirs(state_dir)
|
|
w = _worker(self.db_path,
|
|
executor=self._fake_executor(state_dir))
|
|
self.assertTrue(w._rebuild_host_state_from_checkpoint(
|
|
{"id": "t1", "project_id": "p1"}))
|
|
with open(os.path.join(state_dir, "host-state.json"),
|
|
encoding="utf-8") as fh:
|
|
state = json.load(fh)
|
|
self.assertEqual(state["tasks"]["t1"]["goal_rounds_done"], 2)
|
|
self.assertEqual(state["tasks"]["t1"]["thread_id"], "th1")
|
|
self.assertTrue(state["tasks"]["t1"]["rebuilt_from_checkpoint"])
|
|
|
|
def test_record_step_failure_swallowed(self):
|
|
# 1208-1209: step INSERT raising is logged, never raised.
|
|
w = _worker(self.db_path)
|
|
bad = mock.Mock()
|
|
bad.execute.side_effect = RuntimeError("db down")
|
|
w._record_step(bad, "t1", "round", "summary")
|
|
bad.execute.assert_called()
|
|
|
|
|
|
class WriteResultTestCase(_Base):
|
|
def test_child_pid_branch_writes_full(self):
|
|
# 1664-1670, 1672-1673, 1675: child_pid is not None branch.
|
|
self._insert_task(status="running")
|
|
w = _worker(self.db_path)
|
|
self.assertTrue(w._write_execution_result(
|
|
"t1", result_text="out", result_full="full-out",
|
|
child_pid=1234, turn_id="turn-1"))
|
|
row = self._row()
|
|
self.assertEqual(row["child_pid"], 1234)
|
|
self.assertEqual(row["result_text_full"], "full-out")
|
|
self.assertEqual(row["turn_id"], "turn-1")
|
|
|
|
def test_write_result_non_column_error_reraised(self):
|
|
# 1686: an execute error unrelated to result_text_full is raised.
|
|
self._insert_task(status="running")
|
|
w = _worker(self.db_path)
|
|
w._ensure_result_columns = lambda conn: None
|
|
|
|
class _BadConn:
|
|
def execute(self, *args, **kwargs):
|
|
raise ValueError("boom")
|
|
|
|
def rollback(self):
|
|
pass
|
|
|
|
def commit(self):
|
|
pass
|
|
|
|
def close(self):
|
|
pass
|
|
|
|
with mock.patch("cloud.worker._db_connect",
|
|
return_value=_BadConn()):
|
|
with self.assertRaises(ValueError):
|
|
w._write_execution_result("t1", result_text="out")
|
|
|
|
|
|
class NotifyAndMemoryTestCase(_Base):
|
|
def test_notify_errors_branch_logged(self):
|
|
# 1784-1786: dispatch_event returning errors is logged.
|
|
w = _worker(self.db_path)
|
|
with mock.patch("cloud.push.dispatch_event",
|
|
return_value={"errors": ["sub down"]}):
|
|
w._notify_task_finished("t1", "p1", "completed")
|
|
# No exception, no state change to assert beyond survival.
|
|
|
|
def test_auto_memory_logger_failure_swallowed(self):
|
|
# 1823: logger.warning itself raising inside the injection
|
|
# failure path is swallowed and the task is returned unchanged.
|
|
w = _worker(self.db_path)
|
|
task = {"id": "t1", "project_id": "p1", "created_by": "u1",
|
|
"prompt": "do work"}
|
|
w.logger = mock.Mock()
|
|
w.logger.warning.side_effect = RuntimeError("logger down")
|
|
with mock.patch("cloud.worker.AutoMemoryStore",
|
|
side_effect=RuntimeError("store down")):
|
|
out = w._apply_auto_memory_injection(task)
|
|
self.assertIs(out, task)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|