host
e2d6a454bf
SQLite REAL is 8-byte double, but Postgres REAL is 4-byte float4 (~7 significant digits). Epoch seconds (10 digits) were rounded to 128-second steps, so worker liveness ages were off by up to +-64s and could even go negative (observed live on the host: last_seen ahead of NOW() by 20s). New tables use DOUBLE PRECISION; existing Postgres tables are upgraded only when information_schema actually reports REAL, so the ALTER does not run on every heartbeat (it would take an ACCESS EXCLUSIVE lock each time). SQLite path is unchanged (its REAL is already double precision). Regression test asserts new tables use DOUBLE PRECISION.
484 lines
18 KiB
Python
484 lines
18 KiB
Python
"""多 worker 池派发 + 通用 CloudExecutor + 自动 PR 单测。
|
|
|
|
只用标准库，不跑真实模型回合、不联网、不登录任何服务；
|
|
git/gh 一律用假 runner 注入，不依赖本机是否装 gh、是否有
|
|
远端仓库；CloudExecutor 的 endpoint 一律用假 http_client。
|
|
"""
|
|
|
|
import http.client
|
|
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
|
|
|
|
try:
|
|
from cloud.auth import AuthStore
|
|
from cloud.pr import PrSkipped, create_task_pr
|
|
from cloud.sandbox import CloudExecutor, DockerExecutor
|
|
from cloud.server import create_server, init_db
|
|
from cloud.worker import Worker
|
|
from cloud.workers import (
|
|
heartbeat_worker,
|
|
list_workers,
|
|
register_worker,
|
|
)
|
|
except ImportError:
|
|
sys.path.insert(0, str(Path(__file__).resolve().parents[2]))
|
|
from cloud.auth import AuthStore
|
|
from cloud.pr import PrSkipped, create_task_pr
|
|
from cloud.sandbox import CloudExecutor, DockerExecutor
|
|
from cloud.server import create_server, init_db
|
|
from cloud.worker import Worker
|
|
from cloud.workers import (
|
|
heartbeat_worker,
|
|
list_workers,
|
|
register_worker,
|
|
)
|
|
|
|
|
|
def _result(returncode=0, stdout="", stderr=""):
|
|
return SimpleNamespace(returncode=returncode, stdout=stdout,
|
|
stderr=stderr)
|
|
|
|
|
|
class _ServerCase(unittest.TestCase):
|
|
def setUp(self):
|
|
self.tmp = tempfile.TemporaryDirectory()
|
|
self.db_path = os.path.join(self.tmp.name, "control.db")
|
|
self.auth = AuthStore(self.db_path)
|
|
self.owner_id, self.owner_key = self.auth.create_user("owner")
|
|
self.viewer_id, self.viewer_key = self.auth.create_user("viewer")
|
|
self.server = create_server(self.db_path, "127.0.0.1", 0)
|
|
self.port = self.server.server_address[1]
|
|
self.thread = threading.Thread(
|
|
target=self.server.serve_forever, daemon=True)
|
|
self.thread.start()
|
|
status, self.project = self._req(
|
|
"POST", "/projects", {"name": "p1", "auto_pr": 1},
|
|
self.owner_key)
|
|
self.assertEqual(status, 201, self.project)
|
|
self.auth.set_member(
|
|
self.project["id"], self.owner_id, self.viewer_id, "viewer")
|
|
self.worktrees = os.path.join(self.tmp.name, "worktrees")
|
|
os.makedirs(self.worktrees, exist_ok=True)
|
|
self.server.host_worktrees_dir = self.worktrees
|
|
|
|
def tearDown(self):
|
|
self.server.shutdown()
|
|
self.server.server_close()
|
|
self.thread.join(timeout=5)
|
|
try:
|
|
self.auth.conn.close()
|
|
except Exception:
|
|
pass
|
|
self.tmp.cleanup()
|
|
|
|
def _req(self, method, path, body=None, key=None, headers=None):
|
|
conn = http.client.HTTPConnection(
|
|
"127.0.0.1", self.port, timeout=10)
|
|
hdrs = dict(headers or {})
|
|
payload = None
|
|
if body is not None:
|
|
payload = json.dumps(body).encode("utf-8")
|
|
hdrs.setdefault("Content-Type", "application/json")
|
|
if key is not None:
|
|
hdrs["X-API-Key"] = key
|
|
try:
|
|
conn.request(method, path, body=payload, headers=hdrs)
|
|
resp = conn.getresponse()
|
|
data = resp.read()
|
|
finally:
|
|
conn.close()
|
|
try:
|
|
parsed = json.loads(data.decode("utf-8")) if data else None
|
|
except ValueError:
|
|
parsed = None
|
|
return resp.status, parsed
|
|
|
|
def _insert_task(self, task_id, status, project_id=None):
|
|
conn = sqlite3.connect(self.db_path)
|
|
try:
|
|
conn.execute(
|
|
"INSERT INTO tasks(id, project_id, prompt, status,"
|
|
" depends_on, priority, created_by, created_at,"
|
|
" updated_at, last_heartbeat) VALUES(?, ?, ?, ?, '[]',"
|
|
" 0, 'system', 0, 0, NULL)",
|
|
(task_id, project_id or self.project["id"],
|
|
"do it", status))
|
|
conn.commit()
|
|
finally:
|
|
conn.close()
|
|
|
|
def _task_row(self, task_id):
|
|
conn = sqlite3.connect(self.db_path)
|
|
try:
|
|
row = conn.execute(
|
|
"SELECT status, pr_url, pr_number FROM tasks"
|
|
" WHERE id = ?", (task_id,)).fetchone()
|
|
finally:
|
|
conn.close()
|
|
return row
|
|
|
|
def _audit_actions(self):
|
|
conn = sqlite3.connect(self.db_path)
|
|
try:
|
|
return [r[0] for r in conn.execute(
|
|
"SELECT action FROM audit_log").fetchall()]
|
|
finally:
|
|
conn.close()
|
|
|
|
def _git_runner(self, push_fail=False, remote_fail=False):
|
|
def run(cmd, cwd):
|
|
if "remote" in cmd:
|
|
if remote_fail:
|
|
return _result(1, "", "no remote")
|
|
return _result(0, "https://example.com/repo.git\n")
|
|
if "push" in cmd:
|
|
if push_fail:
|
|
return _result(1, "", "push rejected")
|
|
return _result(0, "")
|
|
return _result(0, "")
|
|
return run
|
|
|
|
def _gh_runner(self, authed=True, pr_url=None):
|
|
def run(cmd, cwd):
|
|
if cmd[:2] == ["gh", "auth"]:
|
|
return _result(0 if authed else 1, "", "")
|
|
if cmd[:2] == ["gh", "pr"]:
|
|
return _result(0, (pr_url or
|
|
"https://example.com/repo/pull/42\n"))
|
|
return _result(0, "")
|
|
return run
|
|
|
|
def _worktree(self, task_id):
|
|
path = os.path.join(self.worktrees, task_id)
|
|
os.makedirs(path, exist_ok=True)
|
|
return path
|
|
|
|
|
|
class WorkerRegistryTestCase(_ServerCase):
|
|
def test_heartbeat_requires_worker_id(self):
|
|
status, body = self._req(
|
|
"POST", "/workers/heartbeat", {}, self.owner_key)
|
|
self.assertEqual(status, 400, body)
|
|
|
|
def test_heartbeat_registers_and_lists_alive(self):
|
|
status, body = self._req(
|
|
"POST", "/workers/heartbeat",
|
|
{"worker_id": "w-docker-1", "executor": "docker",
|
|
"capacity": 2},
|
|
self.owner_key)
|
|
self.assertEqual(status, 200, body)
|
|
self.assertTrue(body["alive"])
|
|
self.assertEqual(body["executor"], "docker")
|
|
status, body = self._req("GET", "/workers", None, self.owner_key)
|
|
self.assertEqual(status, 200, body)
|
|
self.assertEqual(len(body["workers"]), 1)
|
|
self.assertTrue(body["workers"][0]["alive"])
|
|
|
|
def test_heartbeat_preserves_executor_when_not_given(self):
|
|
self._req("POST", "/workers/heartbeat",
|
|
{"worker_id": "w1", "executor": "docker"}, self.owner_key)
|
|
status, body = self._req(
|
|
"POST", "/workers/heartbeat", {"worker_id": "w1"},
|
|
self.owner_key)
|
|
self.assertEqual(status, 200, body)
|
|
self.assertEqual(body["executor"], "docker")
|
|
|
|
def test_stale_worker_lists_not_alive(self):
|
|
register_worker(self.db_path, "w-old", executor="local")
|
|
conn = sqlite3.connect(self.db_path)
|
|
try:
|
|
conn.execute(
|
|
"UPDATE workers SET last_seen = ? WHERE worker_id = ?",
|
|
(time.time() - 3600, "w-old"))
|
|
conn.commit()
|
|
finally:
|
|
conn.close()
|
|
workers = list_workers(self.db_path)
|
|
self.assertEqual(len(workers), 1)
|
|
self.assertFalse(workers[0]["alive"])
|
|
|
|
def test_heartbeat_without_key_is_unauthorized(self):
|
|
status, body = self._req(
|
|
"POST", "/workers/heartbeat", {"worker_id": "w1"}, None)
|
|
self.assertEqual(status, 401, body)
|
|
|
|
|
|
class TaskExecutorFilterTestCase(unittest.TestCase):
|
|
def setUp(self):
|
|
self.tmp = tempfile.TemporaryDirectory()
|
|
init_db(os.path.join(self.tmp.name, "control.db"))
|
|
self.db_path = os.path.join(self.tmp.name, "control.db")
|
|
conn = sqlite3.connect(self.db_path)
|
|
try:
|
|
for tid, exc in (("t-local", "local"), ("t-none", None),
|
|
("t-docker", "docker")):
|
|
conn.execute(
|
|
"INSERT INTO tasks(id, project_id, prompt, status,"
|
|
" depends_on, priority, created_by, created_at,"
|
|
" updated_at, last_heartbeat, executor)"
|
|
" VALUES(?, 'proj', 'do it', 'queued', '[]', 0,"
|
|
" 'system', 0, 0, NULL, ?)", (tid, exc))
|
|
conn.commit()
|
|
finally:
|
|
conn.close()
|
|
|
|
def tearDown(self):
|
|
self.tmp.cleanup()
|
|
|
|
def _docker_worker(self):
|
|
return Worker(
|
|
self.db_path, "host.py", config_path=None, repo=None,
|
|
executor=DockerExecutor(runner=lambda cmd: _result()),
|
|
env={"XCM_API_KEY": "test-placeholder"}, env_key=None,
|
|
task_timeout=1800, max_retries=0, retry_backoff_seconds=0)
|
|
|
|
def test_claim_rejects_other_pool_task(self):
|
|
worker = self._docker_worker()
|
|
self.assertFalse(worker._claim("t-local", executor="docker"))
|
|
conn = sqlite3.connect(self.db_path)
|
|
try:
|
|
status = conn.execute(
|
|
"SELECT status FROM tasks WHERE id = 't-local'"
|
|
).fetchone()[0]
|
|
finally:
|
|
conn.close()
|
|
self.assertEqual(status, "queued")
|
|
|
|
def test_claim_accepts_unassigned_and_same_pool(self):
|
|
worker = self._docker_worker()
|
|
self.assertTrue(worker._claim("t-none", executor="docker"))
|
|
self.assertTrue(worker._claim("t-docker", executor="docker"))
|
|
|
|
def test_cas_race_only_one_worker_wins(self):
|
|
w1 = self._docker_worker()
|
|
w2 = self._docker_worker()
|
|
self.assertTrue(w1._claim("t-none", executor="docker"))
|
|
self.assertFalse(w2._claim("t-none", executor="docker"))
|
|
|
|
def test_pick_next_skips_other_pool(self):
|
|
# 只剩 local 池任务时，docker worker 无就绪任务。
|
|
conn = sqlite3.connect(self.db_path)
|
|
try:
|
|
conn.execute("DELETE FROM tasks WHERE id != 't-local'")
|
|
conn.commit()
|
|
finally:
|
|
conn.close()
|
|
worker = self._docker_worker()
|
|
self.assertIsNone(worker.pick_next())
|
|
|
|
def test_run_one_skips_other_pool_without_model(self):
|
|
conn = sqlite3.connect(self.db_path)
|
|
try:
|
|
conn.execute("DELETE FROM tasks WHERE id != 't-local'")
|
|
conn.commit()
|
|
finally:
|
|
conn.close()
|
|
worker = self._docker_worker()
|
|
result = worker.run_one()
|
|
self.assertFalse(result["dispatched"])
|
|
self.assertEqual(result["reason"], "no_ready_task")
|
|
|
|
|
|
class CloudExecutorTestCase(unittest.TestCase):
|
|
def test_unconfigured_fails_without_fallback(self):
|
|
ex = CloudExecutor(endpoint=None, env={})
|
|
with self.assertRaises(RuntimeError) as ctx:
|
|
ex.execute({"id": "t1", "prompt": "do it"})
|
|
self.assertIn("not configured", str(ctx.exception))
|
|
|
|
def test_configured_posts_only_id_prompt_repo_no_secrets(self):
|
|
seen = {}
|
|
|
|
def fake_client(url, payload, headers):
|
|
seen["url"] = url
|
|
seen["payload"] = payload
|
|
seen["headers"] = headers
|
|
return {"status": "completed", "stdout": "done"}
|
|
|
|
ex = CloudExecutor(
|
|
endpoint="https://runner.example.invalid/run",
|
|
http_client=fake_client,
|
|
auth_header="X-Runner-Token",
|
|
env={"CLOUD_RUNNER_AUTH_TOKEN": "test-token-value"})
|
|
result = ex.execute({"id": "t9", "prompt": "do it",
|
|
"repo": "https://example.com/r.git"})
|
|
self.assertEqual(result["status"], "completed")
|
|
self.assertEqual(seen["url"], "https://runner.example.invalid/run")
|
|
self.assertEqual(set(seen["payload"].keys()),
|
|
{"id", "prompt", "repo_url"})
|
|
self.assertEqual(seen["payload"]["repo_url"],
|
|
"https://example.com/r.git")
|
|
self.assertNotIn("test-token-value",
|
|
json.dumps(seen["payload"]))
|
|
self.assertEqual(seen["headers"]["X-Runner-Token"],
|
|
"test-token-value")
|
|
# dispatch_task 是 worker 走的入口，必须与 execute 同结果。
|
|
result2 = ex.dispatch_task({"id": "t9", "prompt": "do it"})
|
|
self.assertEqual(result2["status"], "completed")
|
|
|
|
def test_failed_status_maps_to_failed(self):
|
|
ex = CloudExecutor(
|
|
endpoint="https://runner.example.invalid/run",
|
|
http_client=lambda u, p, h: {"status": "failed",
|
|
"stdout": "boom"})
|
|
result = ex.execute({"id": "t1", "prompt": "x"})
|
|
self.assertEqual(result["status"], "failed")
|
|
|
|
|
|
class AutoPrTestCase(_ServerCase):
|
|
def _pr(self, task_id, key):
|
|
return self._req("POST", "/tasks/%s/pr" % task_id, {},
|
|
key)
|
|
|
|
def test_auto_pr_disabled_skips_and_stays_merged(self):
|
|
status, proj = self._req(
|
|
"POST", "/projects", {"name": "p-off"}, self.owner_key)
|
|
self.assertEqual(status, 201, proj)
|
|
self._insert_task("toff", "merged", project_id=proj["id"])
|
|
status, body = self._pr("toff", self.owner_key)
|
|
self.assertEqual(status, 200, body)
|
|
self.assertTrue(body["pr_skipped"])
|
|
self.assertEqual(body["reason"], "auto_pr_disabled")
|
|
self.assertEqual(self._task_row("toff")[0], "merged")
|
|
|
|
def test_auto_pr_illegal_value_rejected(self):
|
|
status, body = self._req(
|
|
"POST", "/projects", {"name": "p-bad", "auto_pr": 2},
|
|
self.owner_key)
|
|
self.assertEqual(status, 400, body)
|
|
|
|
def test_non_merged_returns_409(self):
|
|
self._insert_task("tq", "queued")
|
|
status, body = self._pr("tq", self.owner_key)
|
|
self.assertEqual(status, 409, body)
|
|
|
|
def test_viewer_cannot_pr(self):
|
|
self._insert_task("tm", "merged")
|
|
status, body = self._pr("tm", self.viewer_key)
|
|
self.assertEqual(status, 403, body)
|
|
|
|
def test_no_worktree_skips(self):
|
|
self._insert_task("tnw", "merged")
|
|
self.server.pr_git_runner = self._git_runner()
|
|
self.server.pr_gh_runner = self._gh_runner()
|
|
self.server.pr_gh_which = "gh"
|
|
status, body = self._pr("tnw", self.owner_key)
|
|
self.assertEqual(status, 200, body)
|
|
self.assertEqual(body["reason"], "no_worktree")
|
|
self.assertEqual(self._task_row("tnw")[0], "merged")
|
|
|
|
def test_no_remote_skips(self):
|
|
self._insert_task("tnr", "merged")
|
|
self._worktree("tnr")
|
|
self.server.pr_git_runner = self._git_runner(remote_fail=True)
|
|
self.server.pr_gh_runner = self._gh_runner()
|
|
self.server.pr_gh_which = "gh"
|
|
status, body = self._pr("tnr", self.owner_key)
|
|
self.assertEqual(status, 200, body)
|
|
self.assertEqual(body["reason"], "no_remote")
|
|
|
|
def test_no_gh_skips(self):
|
|
self._insert_task("tngh", "merged")
|
|
self._worktree("tngh")
|
|
self.server.pr_git_runner = self._git_runner()
|
|
self.server.pr_gh_runner = self._gh_runner()
|
|
self.server.pr_gh_which = ""
|
|
status, body = self._pr("tngh", self.owner_key)
|
|
self.assertEqual(status, 200, body)
|
|
self.assertEqual(body["reason"], "no_gh")
|
|
|
|
def test_gh_not_authenticated_skips(self):
|
|
self._insert_task("tna", "merged")
|
|
self._worktree("tna")
|
|
self.server.pr_git_runner = self._git_runner()
|
|
self.server.pr_gh_runner = self._gh_runner(authed=False)
|
|
self.server.pr_gh_which = "gh"
|
|
status, body = self._pr("tna", self.owner_key)
|
|
self.assertEqual(status, 200, body)
|
|
self.assertEqual(body["reason"], "gh_not_authenticated")
|
|
|
|
def test_push_failure_returns_502_and_stays_merged(self):
|
|
self._insert_task("tpf", "merged")
|
|
self._worktree("tpf")
|
|
self.server.pr_git_runner = self._git_runner(push_fail=True)
|
|
self.server.pr_gh_runner = self._gh_runner()
|
|
self.server.pr_gh_which = "gh"
|
|
status, body = self._pr("tpf", self.owner_key)
|
|
self.assertEqual(status, 502, body)
|
|
self.assertEqual(body["step"], "push")
|
|
self.assertEqual(self._task_row("tpf")[0], "merged")
|
|
|
|
def test_success_writes_back_pr_and_audit(self):
|
|
self._insert_task("tok", "merged")
|
|
self._worktree("tok")
|
|
self.server.pr_git_runner = self._git_runner()
|
|
self.server.pr_gh_runner = self._gh_runner()
|
|
self.server.pr_gh_which = "gh"
|
|
status, body = self._pr("tok", self.owner_key)
|
|
self.assertEqual(status, 200, body)
|
|
self.assertEqual(body["pr_url"],
|
|
"https://example.com/repo/pull/42")
|
|
self.assertEqual(body["pr_number"], 42)
|
|
self.assertEqual(body["branch"], "host/tok")
|
|
row = self._task_row("tok")
|
|
self.assertEqual(row[0], "merged")
|
|
self.assertEqual(row[1], "https://example.com/repo/pull/42")
|
|
self.assertEqual(row[2], 42)
|
|
self.assertIn("pr_created", self._audit_actions())
|
|
|
|
def test_pr_module_auto_pr_gate(self):
|
|
with self.assertRaises(PrSkipped) as ctx:
|
|
create_task_pr("x", 0)
|
|
self.assertEqual(ctx.exception.reason, "auto_pr_disabled")
|
|
|
|
|
|
class ProjectTaskFieldTestCase(_ServerCase):
|
|
def test_create_task_executor_validation(self):
|
|
status, body = self._req(
|
|
"POST", "/tasks",
|
|
{"project_id": self.project["id"], "prompt": "x",
|
|
"executor": "bogus"},
|
|
self.owner_key)
|
|
self.assertEqual(status, 400, body)
|
|
status, body = self._req(
|
|
"POST", "/tasks",
|
|
{"project_id": self.project["id"], "prompt": "x",
|
|
"executor": "docker"},
|
|
self.owner_key)
|
|
self.assertEqual(status, 201, body)
|
|
self.assertEqual(body["executor"], "docker")
|
|
|
|
|
|
class WorkerTablePrecisionTestCase(unittest.TestCase):
|
|
def test_workers_time_columns_are_double_precision(self):
|
|
# 回归：2026-10-26 主机实测，workers.last_seen 用 REAL 时，
|
|
# Postgres 的 float4 把 epoch 秒舍入到 128 秒一档，存活判断
|
|
# 误差可达 ±64 秒、甚至出现负 age。新表必须是 DOUBLE PRECISION。
|
|
import sqlite3 as _sqlite3
|
|
tmp = tempfile.TemporaryDirectory()
|
|
db_path = os.path.join(tmp.name, "w.db")
|
|
from cloud.workers import ensure_workers_table
|
|
ensure_workers_table(db_path)
|
|
conn = _sqlite3.connect(db_path)
|
|
try:
|
|
rows = conn.execute("PRAGMA table_info(workers)").fetchall()
|
|
finally:
|
|
conn.close()
|
|
tmp.cleanup()
|
|
by_name = {r[1]: r[2] for r in rows}
|
|
self.assertIn("DOUBLE", by_name["last_seen"])
|
|
self.assertIn("DOUBLE", by_name["created_at"])
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|