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

"""goal 字段与 PR 事件推送的端点级单测。

- POST /tasks 接受 goal/verify_command/goal_max_rounds，非法值 400。
- init_db 迁移后 tasks 表含 goal 三列。
- POST /tasks/{id}/pr 成功时发 review-requested，落盘失败时发
  merge-failed（订阅指向本地回环监听，断言真收到）。
git/gh 一律 fake，不跑真实命令、不联网。
"""

import http.client
import http.server
import json
import os
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.push import create_subscription
    from cloud.server import create_server
except ImportError:
    sys.path.insert(0, str(Path(__file__).resolve().parents[2]))
    from cloud.auth import AuthStore
    from cloud.push import create_subscription
    from cloud.server import create_server


def _result(returncode=0, stdout="", stderr=""):
    return SimpleNamespace(returncode=returncode, stdout=stdout,
                           stderr=stderr)


class _Git:
    """假 git：remote/branch/add/commit/push，记录命令。"""

    def __init__(self, remote, branch, push_rc=0):
        self.calls = []
        self.remote = remote
        self.branch = branch
        self.push_rc = push_rc

    def __call__(self, cmd, cwd):
        self.calls.append(list(cmd))
        if "remote" in cmd:
            return _result(0, self.remote + "\n")
        if "branch" in cmd:
            return _result(0, self.branch + "\n")
        if "add" in cmd:
            return _result(0, "")
        if "commit" in cmd:
            return _result(1, "",
                           "nothing to commit, working tree clean")
        if "push" in cmd:
            return _result(self.push_rc, "",
                           "push rejected" if self.push_rc else "")
        return _result(0, "")


class _Gh:
    """假 gh：auth 通过、pr create 回链接。"""

    def __init__(self, pr_url):
        self.calls = []
        self.pr_url = pr_url

    def __call__(self, cmd, cwd):
        self.calls.append(list(cmd))
        if cmd[:2] == ["gh", "auth"]:
            return _result(0, "", "")
        if cmd[:2] == ["gh", "pr"]:
            return _result(0, self.pr_url + "\n")
        return _result(0, "")


class _Listener:
    """本地回环 webhook 接收器。"""

    def __init__(self):
        self.requests = []

        outer = self

        class Handler(http.server.BaseHTTPRequestHandler):
            def do_POST(self):
                length = int(self.headers.get("Content-Length") or 0)
                body = self.rfile.read(length) if length else b""
                outer.requests.append(json.loads(body.decode("utf-8")))
                self.send_response(200)
                self.send_header("Content-Length", "2")
                self.end_headers()
                self.wfile.write(b"{}")

            def log_message(self, *args):
                pass

        self.server = http.server.ThreadingHTTPServer(
            ("127.0.0.1", 0), Handler)
        self.thread = threading.Thread(
            target=self.server.serve_forever, daemon=True)
        self.thread.start()
        self.url = "http://127.0.0.1:%d/hook" % \
            self.server.server_address[1]

    def stop(self):
        self.server.shutdown()
        self.server.server_close()
        self.thread.join(timeout=5)


class GoalPrNotifyTestCase(unittest.TestCase):
    def setUp(self):
        # P1 SSRF 门默认挡 loopback；本地监听器测试显式放行，
        # 生产默认值仍关闭。
        self._old_push_allow = os.environ.get(
            "CODEX_PUSH_ALLOW_PRIVATE_TARGETS")
        os.environ["CODEX_PUSH_ALLOW_PRIVATE_TARGETS"] = "1"
        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()
        self.worktrees = os.path.join(self.tmp.name, "worktrees")
        os.makedirs(self.worktrees, exist_ok=True)
        self.server.host_worktrees_dir = self.worktrees
        self.listener = _Listener()

    def tearDown(self):
        self.listener.stop()
        self.server.shutdown()
        self.server.server_close()
        self.thread.join(timeout=5)
        if getattr(self, "_old_push_allow", None) is None:
            os.environ.pop("CODEX_PUSH_ALLOW_PRIVATE_TARGETS", None)
        else:
            os.environ["CODEX_PUSH_ALLOW_PRIVATE_TARGETS"] = (
                self._old_push_allow)
        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 _project(self, name="p"):
        status, project = self._req(
            "POST", "/projects", {"name": name, "auto_pr": 1},
            self.owner_key)
        self.assertEqual(status, 201, project)
        return project["id"]

    def _merged_task(self, project_id, task_id):
        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(?, ?, ?, 'merged',"
                " '[]', 0, 'system', 0, 0, NULL)",
                (task_id, project_id, "do it"))
            conn.commit()
        finally:
            conn.close()
        os.makedirs(os.path.join(self.worktrees, task_id), exist_ok=True)

    def test_create_task_goal_fields(self):
        pid = self._project()
        status, task = self._req(
            "POST", "/tasks",
            {"project_id": pid, "prompt": "do it",
             "goal": "make x", "verify_command": "test -f x",
             "goal_max_rounds": 2},
            self.owner_key)
        self.assertEqual(status, 201, task)
        self.assertEqual(task["goal"], "make x")
        self.assertEqual(task["verify_command"], "test -f x")
        self.assertEqual(task["goal_max_rounds"], 2)
        # 不带 goal 字段：普通任务，字段为空。
        status, task = self._req(
            "POST", "/tasks",
            {"project_id": pid, "prompt": "plain"},
            self.owner_key)
        self.assertEqual(status, 201, task)
        self.assertIsNone(task["goal"])
        self.assertIsNone(task["verify_command"])
        # 非法轮次：0 与非数字都 400，且不建任务。
        for bad in (0, -1, "x"):
            status, body = self._req(
                "POST", "/tasks",
                {"project_id": pid, "prompt": "bad",
                 "verify_command": "true", "goal_max_rounds": bad},
                self.owner_key)
            self.assertEqual(status, 400, body)

    def test_migration_adds_goal_columns(self):
        conn = sqlite3.connect(self.db_path)
        try:
            cols = {row[1] for row in conn.execute(
                "PRAGMA table_info(tasks)").fetchall()}
        finally:
            conn.close()
        self.assertIn("goal", cols)
        self.assertIn("verify_command", cols)
        self.assertIn("goal_max_rounds", cols)

    def test_pr_success_notifies_review_requested(self):
        pid = self._project()
        create_subscription(
            self.db_path, pid, self.owner_id, "webhook",
            self.listener.url, "instant")
        self._merged_task(pid, "t-pr-ok")
        self.server.pr_git_runner = _Git(
            "https://github.com/octo/repo.git", "host/t-pr-ok")
        self.server.pr_gh_runner = _Gh(
            "https://github.com/octo/repo/pull/9")
        status, body = self._req(
            "POST", "/tasks/t-pr-ok/pr", {}, self.owner_key)
        self.assertEqual(status, 200, body)
        self.assertEqual(len(self.listener.requests), 1)
        event = self.listener.requests[0]
        self.assertEqual(event["event_type"], "review-requested")
        self.assertEqual(event["task_id"], "t-pr-ok")
        self.assertIn("pull/9", event["payload"]["pr_url"])

    def test_pr_failure_notifies_merge_failed(self):
        pid = self._project()
        create_subscription(
            self.db_path, pid, self.owner_id, "webhook",
            self.listener.url, "instant")
        self._merged_task(pid, "t-pr-fail")
        self.server.pr_git_runner = _Git(
            "https://github.com/octo/repo.git", "host/t-pr-fail",
            push_rc=1)
        self.server.pr_gh_runner = _Gh(
            "https://github.com/octo/repo/pull/9")
        status, body = self._req(
            "POST", "/tasks/t-pr-fail/pr", {}, self.owner_key)
        self.assertEqual(status, 502, body)
        self.assertEqual(body["step"], "push")
        self.assertEqual(len(self.listener.requests), 1)
        event = self.listener.requests[0]
        self.assertEqual(event["event_type"], "merge-failed")
        self.assertEqual(event["task_id"], "t-pr-fail")
        self.assertEqual(event["payload"]["step"], "push")


if __name__ == "__main__":
    unittest.main()