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).
550 lines
21 KiB
Python
550 lines
21 KiB
Python
"""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()
|