589 lines
22 KiB
Python
589 lines
22 KiB
Python
"""Coverage batch v4-D2: cloud/worker.py tree/quota/cancel/watchdog/
|
||
cleanup/gateway/process branches.
|
||
|
||
Stubbing discipline: the method under test is never stubbed. Patched
|
||
collaborators are: _tenant_safe_name / _tenant_dirs / _project_repo_url
|
||
/ _quota_limit_and_used / _list_gateway_models / _gateway_name /
|
||
_gateway_origin / ready_tasks / orphaned_containers / CostStore /
|
||
_db_connect / dispatch sinks / instance-level _kill_proc_group /
|
||
_set_child_pid / get_task / heartbeat (data, sink and predicate
|
||
collaborators). Every assertion reads a real return value or a real
|
||
DB row, except watchdog/cleanup tests where the assertion is the
|
||
returned list or the absence of a raise plus a real side-effect set.
|
||
"""
|
||
|
||
import os
|
||
import sqlite3
|
||
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 # 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 _StopOnce:
|
||
"""wait() returns False once, then True."""
|
||
|
||
def __init__(self):
|
||
self.n = 0
|
||
|
||
def wait(self, timeout=None):
|
||
self.n += 1
|
||
return self.n > 1
|
||
|
||
|
||
class TreeSnapshotTestCase(_Base):
|
||
def _docker_worker(self):
|
||
# Subclass of the real DockerExecutor (no-op __init__) so the
|
||
# isinstance check in _tree_snapshot_source passes without
|
||
# patching the module-level class.
|
||
from cloud.worker import DockerExecutor as _RealDocker
|
||
|
||
class _FakeDocker(_RealDocker):
|
||
def __init__(self):
|
||
pass
|
||
|
||
w = _worker(self.db_path, executor=_FakeDocker())
|
||
w.tenants_root = self.tmp.name
|
||
w.tenant_isolation = True
|
||
return w
|
||
|
||
def test_safe_name_failure_returns_none(self):
|
||
# 2003-2004
|
||
w = self._docker_worker()
|
||
w._tenant_safe_name = mock.Mock(side_effect=RuntimeError("bad"))
|
||
self.assertIsNone(w._tree_snapshot_source({"project_id": "p1"}))
|
||
|
||
def test_checkout_exists_returns_checkout(self):
|
||
# 2007
|
||
w = self._docker_worker()
|
||
w._tenant_safe_name = mock.Mock(return_value="safe")
|
||
os.makedirs(os.path.join(self.tmp.name, "safe", "safe"))
|
||
self.assertEqual(
|
||
w._tree_snapshot_source({"project_id": "p1"}),
|
||
os.path.join(self.tmp.name, "safe", "safe"))
|
||
|
||
def test_tenant_dirs_failure_returns_none(self):
|
||
# 2017-2018
|
||
w = self._docker_worker()
|
||
w._tenant_safe_name = mock.Mock(return_value="safe")
|
||
w._project_repo_url = mock.Mock(return_value=None)
|
||
w._tenant_dirs = mock.Mock(side_effect=RuntimeError("dirs down"))
|
||
self.assertIsNone(w._tree_snapshot_source({"project_id": "p1"}))
|
||
|
||
def test_test_repo_missing_returns_none(self):
|
||
# 2023
|
||
w = self._docker_worker()
|
||
w._tenant_safe_name = mock.Mock(return_value="safe")
|
||
w._project_repo_url = mock.Mock(return_value=None)
|
||
w._tenant_dirs = mock.Mock(
|
||
return_value={"test_repo": "/nonexistent/repo-xyz"})
|
||
self.assertIsNone(w._tree_snapshot_source({"project_id": "p1"}))
|
||
|
||
def test_injection_source_failure_swallowed(self):
|
||
# 2051-2053, 2056
|
||
w = _worker(self.db_path)
|
||
w._tree_snapshot_source = mock.Mock(
|
||
side_effect=RuntimeError("snap down"))
|
||
task = {"id": "t1", "project_id": "p1", "prompt": "do work"}
|
||
self.assertIs(w._apply_tree_injection(task), task)
|
||
|
||
def test_injection_logger_failure_swallowed(self):
|
||
# 2054-2055
|
||
w = _worker(self.db_path)
|
||
w._tree_snapshot_source = mock.Mock(
|
||
side_effect=RuntimeError("snap down"))
|
||
w.logger = mock.Mock()
|
||
w.logger.warning.side_effect = RuntimeError("logger down")
|
||
task = {"id": "t1", "project_id": "p1", "prompt": "do work"}
|
||
self.assertIs(w._apply_tree_injection(task), task)
|
||
|
||
|
||
class QuotaTestCase(_Base):
|
||
def test_limit_commit_failure_swallowed_unlimited(self):
|
||
# 2131-2132: commit raising, explicit-NULL quota row -> None.
|
||
conn = sqlite3.connect(self.db_path)
|
||
try:
|
||
conn.execute(
|
||
"INSERT INTO user_quotas(user_id, quota_usd_micros,"
|
||
" updated_at) VALUES(?, NULL, ?)", ("u1", time.time()))
|
||
conn.commit()
|
||
finally:
|
||
conn.close()
|
||
w = _worker(self.db_path)
|
||
real = sqlite3.connect(self.db_path)
|
||
|
||
class _Proxy:
|
||
def __init__(self, conn):
|
||
self._conn = conn
|
||
|
||
def __getattr__(self, name):
|
||
return getattr(self._conn, name)
|
||
|
||
def commit(self):
|
||
raise RuntimeError("commit down")
|
||
|
||
def close(self):
|
||
self._conn.close()
|
||
|
||
with mock.patch("cloud.worker._db_connect",
|
||
return_value=_Proxy(real)):
|
||
quota, used, month, unpriced = w._quota_limit_and_used("u1")
|
||
self.assertIsNone(quota)
|
||
self.assertEqual((used, unpriced), (0, 0))
|
||
|
||
def test_limit_coststore_close_failure_swallowed(self):
|
||
# 2142-2143
|
||
from cloud.cost import CostStore
|
||
CostStore(self.db_path).close() if hasattr(
|
||
CostStore(self.db_path), "close") else None
|
||
store = CostStore(self.db_path)
|
||
try:
|
||
store._conn.close()
|
||
except Exception:
|
||
pass
|
||
w = _worker(self.db_path)
|
||
|
||
class _BadStore:
|
||
def __init__(self, db_path):
|
||
self._conn = mock.Mock()
|
||
self._conn.close.side_effect = RuntimeError("close down")
|
||
|
||
with mock.patch("cloud.worker.CostStore", _BadStore):
|
||
quota, used, month, unpriced = w._quota_limit_and_used("u9")
|
||
self.assertIsNotNone(quota)
|
||
self.assertEqual((used, unpriced), (0, 0))
|
||
|
||
def test_quota_parent_row_type_error(self):
|
||
# 2196-2197
|
||
w = _worker(self.db_path)
|
||
|
||
class _Cursor:
|
||
def fetchone(self):
|
||
return 5
|
||
|
||
class _Conn:
|
||
def execute(self, *a, **k):
|
||
return _Cursor()
|
||
|
||
def close(self):
|
||
pass
|
||
|
||
with mock.patch("cloud.worker._db_connect", return_value=_Conn()):
|
||
msg = w._quota_exceeded({"parent_task_id": "p0"}, None)
|
||
self.assertIn("缺少属主", msg)
|
||
|
||
def test_quota_parent_lookup_failure(self):
|
||
# 2198-2199
|
||
w = _worker(self.db_path)
|
||
with mock.patch("cloud.worker._db_connect",
|
||
side_effect=RuntimeError("db down")):
|
||
msg = w._quota_exceeded({"parent_task_id": "p0"}, None)
|
||
self.assertIn("缺少属主", msg)
|
||
|
||
def test_quota_limit_failure_logger_failure(self):
|
||
# 2213-2214, both on_error flavours
|
||
w = _worker(self.db_path)
|
||
w._quota_limit_and_used = mock.Mock(
|
||
side_effect=RuntimeError("quota db down"))
|
||
w.logger = mock.Mock()
|
||
w.logger.warning.side_effect = RuntimeError("logger down")
|
||
self.assertIn("quota_check_unavailable",
|
||
w._quota_exceeded({"created_by": "u1"}, None,
|
||
on_error="closed"))
|
||
self.assertIn("quota_check_failed",
|
||
w._quota_exceeded({"created_by": "u1"}, None))
|
||
|
||
def test_ready_budget_exceeded_empty(self):
|
||
# 2252
|
||
w = _worker(self.db_path)
|
||
w.usage_turns = 5
|
||
w.max_turns = 1
|
||
self.assertEqual(w._ready_tasks(), [])
|
||
|
||
def test_ready_value_error_empty(self):
|
||
# 2258-2259
|
||
self._insert_task()
|
||
w = _worker(self.db_path)
|
||
with mock.patch("cloud.worker.ready_tasks",
|
||
side_effect=ValueError("cycle")):
|
||
self.assertEqual(w._ready_tasks(), [])
|
||
|
||
|
||
class HeartbeatCancelTestCase(_Base):
|
||
def test_heartbeat_task_failure_swallowed(self):
|
||
# 2330-2331
|
||
w = _worker(self.db_path)
|
||
w._running_set("t1")
|
||
w.heartbeat = mock.Mock(side_effect=RuntimeError("beat down"))
|
||
w._heartbeat_running_tasks()
|
||
|
||
def test_cancel_proc_wait_failure_swallowed(self):
|
||
# 2349-2350
|
||
self._insert_task(status="running")
|
||
w = _worker(self.db_path)
|
||
proc = mock.Mock()
|
||
proc.wait.side_effect = RuntimeError("wait down")
|
||
w._running_set("t1", proc=proc)
|
||
w._kill_proc_group = lambda p: None
|
||
self.assertTrue(w.cancel("t1"))
|
||
self.assertEqual(self._row()["status"], "interrupted")
|
||
|
||
def test_cancel_container_kill_failure_swallowed(self):
|
||
# 2359-2360
|
||
self._insert_task(status="running")
|
||
w = _worker(self.db_path, executor=SimpleNamespace(
|
||
kill_container=mock.Mock(
|
||
side_effect=RuntimeError("kill down"))))
|
||
w._running_set("t1", container="c1")
|
||
self.assertTrue(w.cancel("t1"))
|
||
self.assertEqual(self._row()["status"], "interrupted")
|
||
|
||
|
||
class WatchdogTestCase(_Base):
|
||
def test_cancel_watchdog_lookup_failure_continues(self):
|
||
# 2407-2408
|
||
w = _worker(self.db_path)
|
||
w.get_task = mock.Mock(side_effect=RuntimeError("db down"))
|
||
w._docker_cancel_watchdog("t1", "c1", _StopOnce())
|
||
|
||
def test_cancel_watchdog_kill_failure_swallowed(self):
|
||
# 2420-2421
|
||
w = _worker(self.db_path, executor=SimpleNamespace(
|
||
kill_container=mock.Mock(
|
||
side_effect=RuntimeError("kill down"))))
|
||
w._cancel_requested.add("t1")
|
||
w._docker_cancel_watchdog("t1", "c1", _StopOnce())
|
||
|
||
def test_quota_watchdog_check_failure_continues(self):
|
||
# 2446-2447
|
||
w = _worker(self.db_path)
|
||
w._docker_quota_watchdog(
|
||
"t1", "c1", _StopOnce(),
|
||
lambda: (_ for _ in ()).throw(RuntimeError("check down")))
|
||
|
||
def test_quota_watchdog_kill_failure_swallowed(self):
|
||
# 2459-2460
|
||
w = _worker(self.db_path, executor=SimpleNamespace(
|
||
kill_container=mock.Mock(
|
||
side_effect=RuntimeError("kill down"))))
|
||
w._docker_quota_watchdog("t1", "c1", _StopOnce(), lambda: True)
|
||
self.assertIn("t1", w._quota_hit_tasks)
|
||
|
||
|
||
class CleanupTestCase(_Base):
|
||
def test_orphan_no_executor_caps_returns_empty(self):
|
||
# 2485
|
||
w = _worker(self.db_path)
|
||
self.assertEqual(w.cleanup_orphaned_containers(), [])
|
||
|
||
def test_orphan_list_failure_returns_empty(self):
|
||
# 2488-2491
|
||
w = _worker(self.db_path, executor=SimpleNamespace(
|
||
list_containers=mock.Mock(
|
||
side_effect=RuntimeError("list down")),
|
||
kill_container=mock.Mock()))
|
||
self.assertEqual(w.cleanup_orphaned_containers(), [])
|
||
|
||
def test_orphan_identify_failure_returns_empty(self):
|
||
# 2496-2499
|
||
w = _worker(self.db_path, executor=SimpleNamespace(
|
||
list_containers=mock.Mock(return_value=["c1"]),
|
||
kill_container=mock.Mock()))
|
||
with mock.patch("cloud.worker.orphaned_containers",
|
||
side_effect=RuntimeError("identify down")):
|
||
self.assertEqual(w.cleanup_orphaned_containers(), [])
|
||
|
||
def test_orphan_kill_failure_skipped(self):
|
||
# 2509-2510
|
||
kill = mock.Mock(side_effect=RuntimeError("kill down"))
|
||
w = _worker(self.db_path, executor=SimpleNamespace(
|
||
list_containers=mock.Mock(return_value=["c1"]),
|
||
kill_container=kill))
|
||
with mock.patch("cloud.worker.orphaned_containers",
|
||
return_value=[{"container": "c1",
|
||
"pid_verified": True}]):
|
||
self.assertEqual(w.cleanup_orphaned_containers(), [])
|
||
|
||
def test_stale_cleanup_failure_returns_empty(self):
|
||
# 2534-2537
|
||
w = _worker(self.db_path, executor=SimpleNamespace(
|
||
cleanup_stale_containers=mock.Mock(
|
||
side_effect=RuntimeError("cleanup down"))))
|
||
self.assertEqual(w.cleanup_stale_containers(), [])
|
||
|
||
def test_stale_cleanup_nonempty_returned(self):
|
||
# 2539
|
||
w = _worker(self.db_path, executor=SimpleNamespace(
|
||
cleanup_stale_containers=mock.Mock(return_value=["c1"])))
|
||
self.assertEqual(w.cleanup_stale_containers(), ["c1"])
|
||
|
||
|
||
class GatewayTestCase(_Base):
|
||
def test_key_available_env_failure_false(self):
|
||
# 2564-2565
|
||
class _BadEnv:
|
||
def get(self, *a, **k):
|
||
# 演练注：main 侧 Worker.__init__ 新增路由策略
|
||
# 环境读取（CLOUD_ROUTING_POLICY_FILE/JSON，
|
||
# 未包在 try 内）；假 env 若全键抛错，构造即
|
||
# 炸、测不到 key_available 本身。故仅对路由
|
||
# 策略键回 None（= 未配策略），其余键仍抛错，
|
||
# key_available 的失败回 False 断言意图不变。
|
||
# 另：__init__ 两处裸 env.get 是否应包 try，
|
||
# 属 main 侧跟进项，本演练不改生产 __init__。
|
||
key = a[0] if a else k.get("key")
|
||
if key in ("CLOUD_ROUTING_POLICY_FILE",
|
||
"CLOUD_ROUTING_POLICY_JSON"):
|
||
return None
|
||
raise RuntimeError("env down")
|
||
|
||
w = Worker(self.db_path, "/fake/host.py", env=_BadEnv())
|
||
self.assertFalse(w.key_available())
|
||
|
||
def _sub2api_env(self):
|
||
return (mock.patch("cloud.worker._gateway_name",
|
||
return_value="sub2api"),
|
||
mock.patch("cloud.worker._gateway_origin",
|
||
return_value="http://gw.example"))
|
||
|
||
def test_gateway_parent_row_type_error(self):
|
||
# 2607-2608 -> owner missing -> RuntimeError
|
||
w = _worker(self.db_path)
|
||
name_p, origin_p = self._sub2api_env()
|
||
|
||
class _Cursor:
|
||
def fetchone(self):
|
||
return 5
|
||
|
||
class _Conn:
|
||
def execute(self, *a, **k):
|
||
return _Cursor()
|
||
|
||
def close(self):
|
||
pass
|
||
|
||
with name_p, origin_p, mock.patch(
|
||
"cloud.worker._db_connect", return_value=_Conn()):
|
||
with self.assertRaises(RuntimeError):
|
||
w._gateway_context({"parent_task_id": "p0"})
|
||
|
||
def test_gateway_parent_lookup_failure(self):
|
||
# 2609-2610 -> owner missing -> RuntimeError
|
||
w = _worker(self.db_path)
|
||
name_p, origin_p = self._sub2api_env()
|
||
with name_p, origin_p, mock.patch(
|
||
"cloud.worker._db_connect",
|
||
side_effect=RuntimeError("db down")):
|
||
with self.assertRaises(RuntimeError):
|
||
w._gateway_context({"parent_task_id": "p0"})
|
||
|
||
def test_sub2api_default_model_none_available(self):
|
||
# 2660-2662
|
||
w = _worker(self.db_path)
|
||
with mock.patch.dict(os.environ,
|
||
{"CLOUD_SUB2API_DEFAULT_MODEL": ""}):
|
||
with mock.patch("cloud.worker._gateway_origin",
|
||
return_value="http://gw.example"):
|
||
with mock.patch("cloud.worker._list_gateway_models",
|
||
return_value=[]):
|
||
with self.assertRaises(RuntimeError):
|
||
w._sub2api_default_model(
|
||
{"SUB2API_API_KEY": "k"})
|
||
|
||
def test_sub2api_default_model_first(self):
|
||
# 2663
|
||
w = _worker(self.db_path)
|
||
with mock.patch.dict(os.environ,
|
||
{"CLOUD_SUB2API_DEFAULT_MODEL": ""}):
|
||
with mock.patch("cloud.worker._gateway_origin",
|
||
return_value="http://gw.example"):
|
||
with mock.patch("cloud.worker._list_gateway_models",
|
||
return_value=["m1", "m2"]):
|
||
self.assertEqual(
|
||
w._sub2api_default_model(
|
||
{"SUB2API_API_KEY": "k"}), "m1")
|
||
|
||
def test_task_seed_dir_db_failure_none(self):
|
||
# 2680-2681
|
||
w = _worker(self.db_path)
|
||
with mock.patch("cloud.worker._db_connect",
|
||
side_effect=RuntimeError("db down")):
|
||
self.assertIsNone(w._task_seed_dir({"id": "t1"}))
|
||
|
||
def test_cleanup_seed_failure_swallowed(self):
|
||
# 2701-2702
|
||
w = _worker(self.db_path)
|
||
w._task_seed_dir = mock.Mock(
|
||
side_effect=RuntimeError("seed down"))
|
||
w._cleanup_task_seed({"id": "t1"})
|
||
|
||
|
||
class ProcessTestCase(_Base):
|
||
def _finished_proc(self, returncode=0):
|
||
proc = mock.Mock()
|
||
proc.poll.return_value = returncode
|
||
proc.communicate.return_value = ("out", "err")
|
||
proc.returncode = returncode
|
||
proc.pid = 4242
|
||
return proc
|
||
|
||
def _proc_worker(self, proc):
|
||
return _worker(self.db_path,
|
||
executor=SimpleNamespace(
|
||
popen_task=mock.Mock(return_value=proc)))
|
||
|
||
def test_executor_popen_path(self):
|
||
# 2730 + 2820 (no _result_from_runner on the executor)
|
||
self._insert_task(status="running")
|
||
proc = self._finished_proc()
|
||
w = self._proc_worker(proc)
|
||
out = w._execute_via_process({"id": "t1", "project_id": "p1"})
|
||
self.assertEqual(out["status"], "completed")
|
||
self.assertEqual(out["stdout"], "out")
|
||
|
||
def test_set_child_pid_failure_swallowed(self):
|
||
# 2736-2737
|
||
self._insert_task(status="running")
|
||
w = _worker(self.db_path,
|
||
process_factory=lambda task: self._finished_proc())
|
||
w._set_child_pid = mock.Mock(
|
||
side_effect=RuntimeError("pid down"))
|
||
out = w._execute_via_process({"id": "t1", "project_id": "p1"})
|
||
self.assertEqual(out["status"], "completed")
|
||
|
||
def test_quota_interval_none_defaults(self):
|
||
# 2747
|
||
self._insert_task(status="running")
|
||
w = _worker(self.db_path,
|
||
process_factory=lambda task: self._finished_proc())
|
||
w.quota_poll_interval_seconds = None
|
||
out = w._execute_via_process({"id": "t1", "project_id": "p1"})
|
||
self.assertEqual(out["status"], "completed")
|
||
|
||
def test_loop_db_lookup_failure_continues(self):
|
||
# 2763-2764
|
||
self._insert_task(status="running")
|
||
proc = mock.Mock()
|
||
proc.poll.side_effect = [None, 0]
|
||
proc.communicate.return_value = ("out", "err")
|
||
proc.returncode = 0
|
||
proc.pid = 4242
|
||
w = _worker(self.db_path,
|
||
process_factory=lambda task: proc)
|
||
w.get_task = mock.Mock(side_effect=RuntimeError("db down"))
|
||
out = w._execute_via_process({"id": "t1", "project_id": "p1"})
|
||
self.assertEqual(out["status"], "completed")
|
||
|
||
def test_loop_db_interrupted_cancels(self):
|
||
# 2767-2769
|
||
self._insert_task(status="running")
|
||
proc = mock.Mock()
|
||
proc.poll.side_effect = [None, 0]
|
||
proc.communicate.return_value = ("out", "err")
|
||
proc.returncode = 0
|
||
proc.pid = 4242
|
||
w = _worker(self.db_path,
|
||
process_factory=lambda task: proc)
|
||
w.get_task = mock.Mock(return_value={"status": "interrupted"})
|
||
w._kill_proc_group = lambda p: None
|
||
out = w._execute_via_process({"id": "t1", "project_id": "p1"})
|
||
self.assertEqual(out["status"], "cancelled")
|
||
|
||
def test_loop_quota_check_failure_continues(self):
|
||
# 2780-2781
|
||
self._insert_task(status="running")
|
||
proc = mock.Mock()
|
||
proc.poll.side_effect = [None, 0]
|
||
proc.communicate.return_value = ("out", "err")
|
||
proc.returncode = 0
|
||
proc.pid = 4242
|
||
w = _worker(self.db_path,
|
||
process_factory=lambda task: proc)
|
||
w.quota_poll_interval_seconds = 0
|
||
out = w._execute_via_process(
|
||
{"id": "t1", "project_id": "p1"},
|
||
quota_check=mock.Mock(side_effect=RuntimeError("quota down")))
|
||
self.assertEqual(out["status"], "completed")
|
||
|
||
def test_communicate_failure_falls_back(self):
|
||
# 2789-2794
|
||
self._insert_task(status="running")
|
||
proc = mock.Mock()
|
||
proc.poll.return_value = 0
|
||
proc.communicate.side_effect = [RuntimeError("comm down"),
|
||
("out2", "err2")]
|
||
proc.returncode = 0
|
||
proc.pid = 4242
|
||
w = _worker(self.db_path,
|
||
process_factory=lambda task: proc)
|
||
out = w._execute_via_process({"id": "t1", "project_id": "p1"})
|
||
self.assertEqual(out["stdout"], "out2")
|
||
proc.kill.assert_called_once()
|
||
|
||
def test_goal_result_from_runner(self):
|
||
# 2817
|
||
self._insert_task(status="running")
|
||
proc = self._finished_proc()
|
||
marker = {"status": "goal-met", "task_id": "t1"}
|
||
w = _worker(self.db_path, executor=SimpleNamespace(
|
||
popen_task=mock.Mock(return_value=proc),
|
||
_goal_result_from_runner=mock.Mock(return_value=marker)))
|
||
out = w._execute_via_process({"id": "t1", "project_id": "p1"},
|
||
kind="goal")
|
||
self.assertEqual(out, marker)
|
||
|
||
|
||
if __name__ == "__main__":
|
||
unittest.main()
|