codex-app-server/cloud/tests/test_pr_jobs.py

"""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,
                               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()


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()
        self.assertEqual(worker._run_pr_jobs(), [])

    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())