"""Coverage batch v2: cloud/server.py handler branches. Covers the baseline gaps in _full_auto_branch_protection_ok, _full_auto_github_request, _admin_cost_by_creator, _workers_heartbeat, _admin_backup and _task_trace_stream. Stubbing discipline: the method under test is never stubbed. Its sinks and predicate collaborators are: _send_json (captured output), _read_json (request body), _is_global_admin via the real CODEX_GLOBAL_ADMINS env, _user_exists via a real users row, _cost_by_creator_payload (fixed payload / ValueError for the month branch), _registry_heartbeat (ValueError case only), backup_postgres (a real temp file, pg_dump itself not run), _fetch_task / _trace_snapshot / _auth_store for the SSE stream (the stream loop and emit-failure logic under test stays real). """ import io import os import sqlite3 import sys import tempfile import unittest import urllib.error 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)) import cloud.server as server_mod # noqa: E402 from cloud.handler_fullauto import FullAutoError, _full_auto_branch_protection_ok, _full_auto_github_request # noqa: E402 from cloud.server import ( # noqa: E402 CloudControlPlaneHandler, init_db, ) def _handler(db_path): h = CloudControlPlaneHandler.__new__(CloudControlPlaneHandler) h.server = SimpleNamespace(db_path=db_path) sent = [] h._send_json = lambda status, obj, extra_headers=None: sent.append( (status, obj)) h.sent = sent return h def _task_row(status="running", task_id="T1"): return { "id": task_id, "project_id": "p1", "prompt": "do work", "status": status, "idempotency_key": None, "depends_on": "[]", "priority": 0, "created_by": "tester", "created_at": 0.0, "updated_at": 0.0, "last_heartbeat": None, "result_text": None, "result_text_full": None, "error_text": None, "actual_model": None, "started_at": None, "finished_at": None, } class FullAutoHelpersTestCase(unittest.TestCase): def _valid_protection(self, checks=("verify",)): return { "required_status_checks": {"strict": True, "contexts": list(checks)}, "enforce_admins": True, "required_pull_request_reviews": { "required_approving_review_count": 1}, "allow_force_pushes": False, "allow_deletions": False, } def test_branch_protection_ok(self): self.assertTrue(_full_auto_branch_protection_ok( self._valid_protection(), ["verify"])) self.assertTrue(_full_auto_branch_protection_ok( self._valid_protection(("verify", "lint")), ["verify"])) def test_branch_protection_failures(self): self.assertFalse(_full_auto_branch_protection_ok( "not-a-dict", ["verify"])) self.assertFalse(_full_auto_branch_protection_ok( self._valid_protection(), [])) bad = self._valid_protection() bad["required_status_checks"] = {"strict": False, "contexts": ["verify"]} self.assertFalse(_full_auto_branch_protection_ok(bad, ["verify"])) bad = self._valid_protection() bad["required_status_checks"] = ["verify"] self.assertFalse(_full_auto_branch_protection_ok(bad, ["verify"])) bad = self._valid_protection() bad["required_status_checks"] = {"strict": True, "contexts": ["other"]} self.assertFalse(_full_auto_branch_protection_ok(bad, ["verify"])) bad = self._valid_protection() bad["enforce_admins"] = False self.assertFalse(_full_auto_branch_protection_ok(bad, ["verify"])) bad = self._valid_protection() bad["required_pull_request_reviews"] = None self.assertFalse(_full_auto_branch_protection_ok(bad, ["verify"])) bad = self._valid_protection() bad["allow_force_pushes"] = True self.assertFalse(_full_auto_branch_protection_ok(bad, ["verify"])) bad = self._valid_protection() bad["allow_deletions"] = True self.assertFalse(_full_auto_branch_protection_ok(bad, ["verify"])) def test_github_request_fake_client(self): def fake(method, url, headers, body): return 201, {"ok": True} self.assertEqual(_full_auto_github_request( fake, "POST", "https://example.invalid/x", {}, None), (201, {"ok": True})) with self.assertRaises(FullAutoError): _full_auto_github_request( lambda m, u, h, b: {"not": "a-tuple"}, "GET", "https://example.invalid/x", {}, None) def test_github_request_urllib_paths(self): class _Resp: def __init__(self, raw, status=200): self._raw = raw self.status = status def read(self): return self._raw def __enter__(self): return self def __exit__(self, *args): return False with mock.patch("urllib.request.urlopen", return_value=_Resp(b'{"a": 1}', 200)): self.assertEqual(_full_auto_github_request( None, "GET", "https://example.invalid/x", {}, None), (200, {"a": 1})) err = urllib.error.HTTPError( "https://example.invalid/x", 404, "Not Found", {}, io.BytesIO(b'{"error": "nope"}')) with mock.patch("urllib.request.urlopen", side_effect=err): self.assertEqual(_full_auto_github_request( None, "GET", "https://example.invalid/x", {}, None), (404, {"error": "nope"})) with mock.patch("urllib.request.urlopen", return_value=_Resp(b'not json', 200)): self.assertEqual(_full_auto_github_request( None, "GET", "https://example.invalid/x", {}, None), (200, {})) with mock.patch("urllib.request.urlopen", return_value=_Resp(b'[1, 2]', 200)): self.assertEqual(_full_auto_github_request( None, "GET", "https://example.invalid/x", {}, None), (200, {})) class AdminCostTestCase(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) conn = sqlite3.connect(self.db_path) try: conn.execute( "CREATE TABLE IF NOT EXISTS users(id TEXT PRIMARY KEY," " username TEXT UNIQUE, api_key_hash TEXT," " created_at REAL)") conn.execute( "INSERT OR IGNORE INTO users(id, username, api_key_hash," " created_at) VALUES('u1', 'u1', '', 0)") conn.commit() finally: conn.close() def tearDown(self): self.tmp.cleanup() def _handler(self): h = _handler(self.db_path) audits = [] h._audit_admin = lambda u, a, d: audits.append((u, a, d)) h.audits = audits return h def test_forbidden_for_non_admin(self): h = self._handler() with mock.patch.dict(os.environ, {"CODEX_GLOBAL_ADMINS": "admin1"}, clear=False): h._admin_cost_by_creator("nobody", {"user_id": ["u1"]}) self.assertEqual(h.sent[0][0], 403) self.assertTrue(any(a == "admin_quota_denied" for _u, a, _d in h.audits)) def test_target_validation(self): h = self._handler() with mock.patch.dict(os.environ, {"CODEX_GLOBAL_ADMINS": "admin1"}): h._admin_cost_by_creator("admin1", {}) self.assertEqual(h.sent[0][0], 400) h.sent.clear() h._admin_cost_by_creator("admin1", {"user_id": ["ghost"]}) self.assertEqual(h.sent[0][0], 404) def test_month_value_error_and_success(self): h = self._handler() h._cost_by_creator_payload = mock.Mock( side_effect=ValueError("bad month")) with mock.patch.dict(os.environ, {"CODEX_GLOBAL_ADMINS": "admin1"}): h._admin_cost_by_creator( "admin1", {"user_id": ["u1"], "month": ["2026-13"]}) self.assertEqual(h.sent[0][0], 400) self.assertIn("YYYY-MM", h.sent[0][1]["error"]) h.sent.clear() h._cost_by_creator_payload = mock.Mock( return_value={"user_id": "u1", "month": "2026-11"}) with mock.patch.dict(os.environ, {"CODEX_GLOBAL_ADMINS": "admin1"}): h._admin_cost_by_creator( "admin1", {"user_id": ["u1"], "month": ["2026-11"]}) self.assertEqual(h.sent[0][0], 200) self.assertEqual(h.sent[0][1]["user_id"], "u1") class WorkersHeartbeatTestCase(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 _call(self, body=None, err=None): h = _handler(self.db_path) h._read_json = lambda: (body, err) h._workers_heartbeat("u1") return h.sent def test_read_error(self): sent = self._call(body=None, err="invalid json") self.assertEqual(sent[0][0], 400) def test_validation_errors(self): self.assertEqual(self._call({})[0][1]["error"], "worker_id required") self.assertEqual(self._call({"worker_id": "w1", "capacity": "many"})[0][1]["error"], "capacity must be an integer") self.assertEqual(self._call({"worker_id": "w1", "running": "many"})[0][1]["error"], "running must be an integer") self.assertEqual(self._call({"worker_id": "w1", "running": -1})[0][1]["error"], "running must be >= 0") self.assertEqual(self._call({"worker_id": "w1", "host": 123})[0][1]["error"], "host must be a string") self.assertEqual(self._call({"worker_id": "w1", "pid": "12"})[0][1]["error"], "pid must be an integer") self.assertEqual(self._call({"worker_id": "w1", "pid": 0})[0][1]["error"], "pid must be > 0") self.assertEqual(self._call({"worker_id": "w1", "address": 5})[0][1]["error"], "address must be a string") def test_registry_value_error(self): h = _handler(self.db_path) h._read_json = lambda: ({"worker_id": "w1"}, None) with mock.patch("cloud.server._registry_heartbeat", side_effect=ValueError("worker_id required")): h._workers_heartbeat("u1") self.assertEqual(h.sent[0][0], 400) self.assertIn("worker_id required", h.sent[0][1]["error"]) def test_success_registers(self): sent = self._call({"worker_id": "w1", "capacity": 2, "running": 0, "host": "h1", "pid": 123, "address": "10.0.0.1"}) self.assertEqual(sent[0][0], 200) self.assertTrue(sent[0][1]["alive"]) self.assertEqual(sent[0][1]["worker_id"], "w1") class AdminBackupTestCase(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) self.out_root = os.path.join(self.tmp.name, "backups") os.makedirs(self.out_root) def tearDown(self): self.tmp.cleanup() def _handler(self): h = _handler(self.db_path) h._audit_admin = lambda u, a, d: None return h def test_read_error_and_out_dir_type(self): h = self._handler() h._read_json = lambda: (None, "invalid json") with mock.patch.dict(os.environ, {"CODEX_GLOBAL_ADMINS": "admin1"}): h._admin_backup("admin1") self.assertEqual(h.sent[0][0], 400) h.sent.clear() h._read_json = lambda: ({"out_dir": 123}, None) with mock.patch.dict(os.environ, {"CODEX_GLOBAL_ADMINS": "admin1"}): h._admin_backup("admin1") self.assertEqual(h.sent[0][0], 400) self.assertIn("out_dir must be a string", h.sent[0][1]["error"]) def test_forbidden_and_outside_root(self): h = self._handler() h._read_json = lambda: ({"out_dir": self.out_root}, None) with mock.patch.dict(os.environ, {"CODEX_GLOBAL_ADMINS": "admin1"}): h._admin_backup("nobody") self.assertEqual(h.sent[0][0], 403) h.sent.clear() h._read_json = lambda: ({"out_dir": "/etc"}, None) with mock.patch.dict(os.environ, {"CODEX_GLOBAL_ADMINS": "admin1", "CLOUD_BACKUP_ROOT": self.out_root}): h._admin_backup("admin1") self.assertEqual(h.sent[0][0], 400) def test_sqlite_backup_with_fallback_root(self): # No out_dir and no CLOUD_BACKUP_ROOT: the CLOUD_BACKUP_DIR # fallback root supplies roots[0] as the candidate. h = self._handler() h._read_json = lambda: ({}, None) env = {"CODEX_GLOBAL_ADMINS": "admin1", "CLOUD_BACKUP_DIR": self.out_root} with mock.patch.dict(os.environ, env): os.environ.pop("CLOUD_BACKUP_ROOT", None) os.environ.pop("CLOUD_DB_BACKEND", None) h._admin_backup("admin1") self.assertEqual(h.sent[0][0], 200, h.sent) self.assertTrue(h.sent[0][1]["path"].startswith(self.out_root)) def test_postgres_backup_branch_with_fake_dump(self): dest = os.path.join(self.out_root, "pg.dump") Path(dest).write_bytes(b"fake-dump") h = self._handler() # _db() resolves the backend from env and refuses PG without # a driver; the backup branch under test is the selection + # success response, so the count query runs on the real # SQLite file through a plain connection here. h._db = lambda: sqlite3.connect(self.db_path) h._read_json = lambda: ({"out_dir": self.out_root}, None) env = {"CODEX_GLOBAL_ADMINS": "admin1", "CLOUD_BACKUP_ROOT": self.out_root, "CLOUD_DB_BACKEND": "postgres"} import cloud.backup as backup_mod with mock.patch.dict(os.environ, env): with mock.patch.object(backup_mod, "backup_postgres", return_value=dest): h._admin_backup("admin1") self.assertEqual(h.sent[0][0], 200, h.sent) self.assertEqual(h.sent[0][1]["backend"], "postgres") self.assertEqual(h.sent[0][1]["path"], dest) def test_chmod_failure_swallowed_and_backup_failure_500(self): h = self._handler() h._read_json = lambda: ({"out_dir": self.out_root}, None) env = {"CODEX_GLOBAL_ADMINS": "admin1", "CLOUD_BACKUP_ROOT": self.out_root} with mock.patch.dict(os.environ, env): with mock.patch("os.chmod", side_effect=OSError("chmod down")): h._admin_backup("admin1") self.assertEqual(h.sent[0][0], 200, h.sent) h.sent.clear() import cloud.backup as backup_mod with mock.patch.dict(os.environ, env): with mock.patch.object(backup_mod, "backup_db", side_effect=RuntimeError("dump down")): h._admin_backup("admin1") self.assertEqual(h.sent[0][0], 500) self.assertIn("backup failed", h.sent[0][1]["error"]) class _AuthCtx: def __enter__(self): return SimpleNamespace(check_permission=lambda *args: True) def __exit__(self, *args): return False class _Writer: def __init__(self, fail_after=None): self.chunks = [] self.calls = 0 self.fail_after = fail_after def write(self, data): self.calls += 1 if self.fail_after is not None and self.calls > self.fail_after: raise BrokenPipeError("pipe closed") self.chunks.append(data.decode("utf-8") if isinstance(data, bytes) else data) def flush(self): pass def _trace_handler(db_path, fetch, entries, writer=None, headers_raise=False): h = _handler(db_path) h._fetch_task = fetch h._trace_snapshot = lambda task: (entries, "") h.send_response = lambda *a: (_ for _ in ()).throw( BrokenPipeError("gone")) if headers_raise else None h.send_header = lambda *a: None h.end_headers = lambda: None h.wfile = writer if writer is not None else _Writer() return h class TraceStreamTestCase(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 _run(self, handler): with mock.patch("cloud.server._auth_store", lambda db_path: _AuthCtx()): handler._task_trace_stream("u1", "T1") return "".join(handler.wfile.chunks) def test_headers_broken_pipe_returns(self): h = _trace_handler(self.db_path, lambda tid: _task_row(), [{"id": "e1"}], headers_raise=True) out = self._run(h) self.assertEqual(out, "") def test_snapshot_emit_failure_returns(self): writer = _Writer(fail_after=0) h = _trace_handler(self.db_path, lambda tid: _task_row(), [{"id": "e1"}], writer=writer) out = self._run(h) self.assertEqual(out, "") self.assertEqual(writer.calls, 1) def test_poll_timeout(self): h = _trace_handler(self.db_path, lambda tid: _task_row(), [{"id": "e1"}]) with mock.patch.object(server_mod, "TRACE_STREAM_MAX_SECONDS", 0.05), \ mock.patch.object(server_mod, "TRACE_STREAM_POLL_SECONDS", 0.01): out = self._run(h) self.assertIn('"type": "snapshot"', out) self.assertIn('"reason": "timeout"', out) def test_poll_deleted(self): calls = [0] def _fetch(tid): calls[0] += 1 return _task_row() if calls[0] == 1 else None h = _trace_handler(self.db_path, _fetch, [{"id": "e1"}]) with mock.patch.object(server_mod, "TRACE_STREAM_MAX_SECONDS", 60), \ mock.patch.object(server_mod, "TRACE_STREAM_POLL_SECONDS", 0.01): out = self._run(h) self.assertIn('"status": "deleted"', out) def test_poll_new_entry_then_terminal(self): calls = [0] def _fetch(tid): calls[0] += 1 if calls[0] == 1: return _task_row(status="running") return _task_row(status="completed") def _snapshot(task): if task["status"] == "completed": return [{"id": "e1"}, {"id": "e2"}], "" return [{"id": "e1"}], "" h = _trace_handler(self.db_path, _fetch, [{"id": "e1"}]) h._trace_snapshot = _snapshot with mock.patch.object(server_mod, "TRACE_STREAM_MAX_SECONDS", 60), \ mock.patch.object(server_mod, "TRACE_STREAM_POLL_SECONDS", 0.01): out = self._run(h) self.assertIn('"type": "entry"', out) self.assertIn('"status": "completed"', out) def test_poll_emit_failure_returns(self): writer = _Writer(fail_after=1) calls = [0] def _fetch(tid): calls[0] += 1 if calls[0] == 1: return _task_row(status="running") return _task_row(status="completed") def _snapshot(task): if task["status"] == "completed": return [{"id": "e1"}, {"id": "e2"}], "" return [{"id": "e1"}], "" h = _trace_handler(self.db_path, _fetch, [{"id": "e1"}], writer=writer) h._trace_snapshot = _snapshot with mock.patch.object(server_mod, "TRACE_STREAM_MAX_SECONDS", 60), \ mock.patch.object(server_mod, "TRACE_STREAM_POLL_SECONDS", 0.01): out = self._run(h) # Snapshot write ok, loop entry write fails -> stream ends # without a done event. self.assertIn('"type": "snapshot"', out) self.assertNotIn('"type": "done"', out) def test_poll_fetch_broken_pipe_returns(self): calls = [0] def _fetch(tid): calls[0] += 1 if calls[0] == 1: return _task_row(status="running") raise BrokenPipeError("gone") h = _trace_handler(self.db_path, _fetch, [{"id": "e1"}]) with mock.patch.object(server_mod, "TRACE_STREAM_MAX_SECONDS", 60), \ mock.patch.object(server_mod, "TRACE_STREAM_POLL_SECONDS", 0.01): out = self._run(h) self.assertIn('"type": "snapshot"', out) self.assertNotIn('"reason": "timeout"', out) if __name__ == "__main__": unittest.main()