"""HA 阶段 1 单测：pr 任务队列 + worktree_host + /pr 跨主机分支。 覆盖：pr_jobs 的 enqueue/claim(CAS)/finish 与去重、tasks 的 worktree_host 迁移（server 与 worker 两支）、认领时写主机名、 /pr 在 worktree_host 为别机时排队（不直接跑、不重复排）、 local worker 的 pr 任务循环（CAS 认领 -> 执行 -> 写回终态）、 autoscaler 给 worker 强制追加 --host-name。 git/GitHub 一律 fake 或 monkeypatch，不跑真实命令、不联网。 """ import http.client import json import os import socket import sqlite3 import sys import tempfile import threading import unittest from pathlib import Path from types import SimpleNamespace try: from cloud.auth import AuthStore from cloud.autoscaler import Autoscaler from cloud.db import connect as _db_connect from cloud.pr_exec import (claim_pr_job, enqueue_pr_job, ensure_pr_jobs_table, find_open_pr_job, finish_pr_job) from cloud.server import create_server, migrate_tasks_table from cloud.worker import Worker except ImportError: sys.path.insert(0, str(Path(__file__).resolve().parents[2])) from cloud.auth import AuthStore from cloud.autoscaler import Autoscaler from cloud.db import connect as _db_connect from cloud.pr_exec import (claim_pr_job, enqueue_pr_job, ensure_pr_jobs_table, find_open_pr_job, finish_pr_job) from cloud.server import create_server, migrate_tasks_table from cloud.worker import Worker def _old_tasks_table(db_path): conn = sqlite3.connect(db_path) try: conn.execute( "CREATE TABLE tasks(id TEXT PRIMARY KEY, project_id TEXT," " prompt TEXT, status TEXT, depends_on TEXT, priority INTEGER," " created_by TEXT, created_at REAL, updated_at REAL," " last_heartbeat REAL)") conn.commit() finally: conn.close() def _merged_task(db_path, task_id, project_id, host=None): conn = sqlite3.connect(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(?, ?, ?, 'merged'," " '[]', 0, 'system', 0, 0, NULL)", (task_id, project_id, "do it")) if host is not None: conn.execute("UPDATE tasks SET worktree_host = ? WHERE id = ?", (host, task_id)) conn.commit() finally: conn.close() class PrJobsTableTestCase(unittest.TestCase): def setUp(self): self.tmp = tempfile.TemporaryDirectory() self.db_path = os.path.join(self.tmp.name, "control.db") def tearDown(self): self.tmp.cleanup() def test_enqueue_claim_finish_and_cas(self): conn = _db_connect(self.db_path) try: job1 = enqueue_pr_job(conn, "t1", "p1", "host-a") job2 = enqueue_pr_job(conn, "t2", "p1", "host-a") self.assertIsNotNone(find_open_pr_job(conn, "t1")) claimed = claim_pr_job(conn, "host-a") self.assertEqual(claimed["task_id"], "t1") # 第二次认领拿不到同一任务（CAS），只能拿到下一个。 claimed2 = claim_pr_job(conn, "host-a") self.assertEqual(claimed2["task_id"], "t2") self.assertIsNone(claim_pr_job(conn, "host-a")) # 别的主机认领不到 host-a 的任务。 self.assertIsNone(claim_pr_job(conn, "host-b")) finish_pr_job(conn, job1, "completed", "https://github.com/o/a/pull/1", 1) self.assertIsNone(find_open_pr_job(conn, "t1")) row = conn.execute( "SELECT status, pr_url, pr_number FROM pr_jobs" " WHERE id = ?", (job1,)).fetchone() self.assertEqual(row[0], "completed") self.assertEqual(row[1], "https://github.com/o/a/pull/1") self.assertEqual(row[2], 1) finally: conn.close() # -- G13：队列记费用归属人 + 老表补列 -------------------------------- def test_enqueue_records_created_by(self): conn = _db_connect(self.db_path) try: enqueue_pr_job(conn, "t1", "p1", "host-a", created_by="user-9") job = find_open_pr_job(conn, "t1") self.assertEqual(job["created_by"], "user-9") # 不传时为 None（老调用方），worker 回落 system。 enqueue_pr_job(conn, "t2", "p1", "host-a") self.assertIsNone(find_open_pr_job(conn, "t2")["created_by"]) finally: conn.close() def test_old_table_migration_adds_created_by(self): conn = sqlite3.connect(self.db_path) try: conn.execute( "CREATE TABLE pr_jobs(id TEXT PRIMARY KEY," " task_id TEXT NOT NULL, project_id TEXT NOT NULL," " host TEXT NOT NULL, status TEXT NOT NULL, pr_url TEXT," " pr_number INTEGER, error TEXT, created_at REAL," " updated_at REAL, started_at REAL, finished_at REAL)") conn.commit() finally: conn.close() conn = _db_connect(self.db_path) try: ensure_pr_jobs_table(conn) cols = {row[1] for row in conn.execute( "PRAGMA table_info(pr_jobs)").fetchall()} self.assertIn("created_by", cols) # 补列后照常排队并记归属人。 enqueue_pr_job(conn, "t1", "p1", "host-a", created_by="user-9") self.assertEqual( find_open_pr_job(conn, "t1")["created_by"], "user-9") finally: conn.close() class WorktreeHostMigrationTestCase(unittest.TestCase): def setUp(self): self.tmp = tempfile.TemporaryDirectory() self.db_path = os.path.join(self.tmp.name, "control.db") _old_tasks_table(self.db_path) def tearDown(self): self.tmp.cleanup() def test_server_migrate_adds_worktree_host(self): conn = _db_connect(self.db_path) try: migrate_tasks_table(conn) cols = {row[1] for row in conn.execute( "PRAGMA table_info(tasks)").fetchall()} self.assertIn("worktree_host", cols) # 再跑一次是 no-op。 migrate_tasks_table(conn) finally: conn.close() def test_worker_ensure_adds_worktree_host(self): worker = Worker(self.db_path, "/nonexistent/host.py", env_key="TEST_KEY") conn = _db_connect(self.db_path) try: worker._ensure_result_columns(conn) cols = {row[1] for row in conn.execute( "PRAGMA table_info(tasks)").fetchall()} self.assertIn("worktree_host", cols) finally: conn.close() def test_claim_writes_host_name(self): worker = Worker(self.db_path, "/nonexistent/host.py", env_key="TEST_KEY", host_name="worker-host-9") conn = _db_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('t1', 'p1', 'x'," " 'queued', '[]', 0, 'system', 0, 0, NULL)") conn.commit() finally: conn.close() self.assertTrue(worker._claim("t1", executor="local")) conn = _db_connect(self.db_path) try: row = conn.execute( "SELECT status, worktree_host FROM tasks WHERE id = 't1'" ).fetchone() self.assertEqual(row[0], "running") self.assertEqual(row[1], "worker-host-9") finally: conn.close() class PrEndpointRemoteTestCase(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.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() 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): conn = http.client.HTTPConnection( "127.0.0.1", self.port, timeout=10) hdrs = {} payload = None if body is not None: payload = json.dumps(body).encode("utf-8") hdrs["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 test_remote_host_queues_and_dedupes(self): status, project = self._req( "POST", "/projects", {"name": "p", "auto_pr": 1}, self.owner_key) self.assertEqual(status, 201, project) # 项目成员默认 owner 即 maintainer（create 时建成员行）。 _merged_task(self.db_path, "t-remote", project["id"], host="elsewhere-host-1") status, body = self._req( "POST", "/tasks/t-remote/pr", {}, self.owner_key) self.assertEqual(status, 200, body) self.assertTrue(body.get("pr_queued")) self.assertEqual(body.get("worktree_host"), "elsewhere-host-1") job_id = body.get("pr_job_id") self.assertIsNotNone(job_id) # 再调一次：去重，返回同一个 job，不重复排队。 status, body2 = self._req( "POST", "/tasks/t-remote/pr", {}, self.owner_key) self.assertEqual(status, 200, body2) self.assertEqual(body2.get("pr_job_id"), job_id) conn = _db_connect(self.db_path) try: rows = conn.execute( "SELECT COUNT(*) FROM pr_jobs WHERE task_id = 't-remote'" ).fetchone() self.assertEqual(rows[0], 1) finally: conn.close() class WorkerPrJobsTestCase(unittest.TestCase): def setUp(self): self.tmp = tempfile.TemporaryDirectory() self.db_path = os.path.join(self.tmp.name, "control.db") _old_tasks_table(self.db_path) def tearDown(self): self.tmp.cleanup() def test_local_worker_runs_pr_job(self): import cloud.worker as worker_mod calls = [] def _fake_execute(db_path, task_id, user_id, **kwargs): calls.append((task_id, user_id)) return {"pr_url": "https://github.com/o/a/pull/9", "pr_number": 9, "branch": "host/t1", "auth_via": "pat"} worker = Worker(self.db_path, "/nonexistent/host.py", env_key="TEST_KEY", host_name="elsewhere-host-1") conn = _db_connect(self.db_path) try: worker._ensure_result_columns(conn) conn.execute( "INSERT INTO tasks(id, project_id, prompt, status," " depends_on, priority, created_by, created_at," " updated_at, last_heartbeat) VALUES('t1', 'p1', 'x'," " 'merged', '[]', 0, 'system', 0, 0, NULL)") conn.commit() enqueue_pr_job(conn, "t1", "p1", "elsewhere-host-1") enqueue_pr_job(conn, "t2", "p1", "other-host") finally: conn.close() orig = worker_mod.execute_task_pr worker_mod.execute_task_pr = _fake_execute try: results = worker._run_pr_jobs() finally: worker_mod.execute_task_pr = orig self.assertEqual(len(results), 1) self.assertEqual(results[0]["task_id"], "t1") self.assertEqual(results[0]["status"], "completed") self.assertEqual(calls, [("t1", "system")]) conn = _db_connect(self.db_path) try: row = conn.execute( "SELECT status, pr_url FROM pr_jobs WHERE task_id = 't1'" ).fetchone() self.assertEqual(row[0], "completed") self.assertEqual(row[1], "https://github.com/o/a/pull/9") # 别的主机的任务没被碰。 row2 = conn.execute( "SELECT status FROM pr_jobs WHERE task_id = 't2'" ).fetchone() self.assertEqual(row2[0], "queued") finally: conn.close() def test_docker_worker_skips_pr_jobs(self): from cloud.sandbox import DockerExecutor worker = Worker(self.db_path, "/nonexistent/host.py", env_key="TEST_KEY", executor=DockerExecutor("/fake/host.py")) conn = _db_connect(self.db_path) try: enqueue_pr_job(conn, "t1", "p1", worker.host_name) finally: conn.close() # P1：Docker worker 不再静默返回空，改为显式 needs_local_worker， # 行状态持久化为 needs_local_worker、仍可被 local worker 认领。 jobs = worker._run_pr_jobs() self.assertEqual(len(jobs), 1) self.assertEqual(jobs[0]["status"], "needs_local_worker") self.assertEqual(jobs[0]["task_id"], "t1") def test_autoscaler_forces_host_name(self): scaler = Autoscaler(self.db_path, []) cmd = scaler._worker_command("w1") self.assertIn("--host-name", cmd) self.assertEqual(cmd[cmd.index("--host-name") + 1], socket.gethostname())