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

"""控制面 server 集成测试：对本地临时端口做真实 HTTP 请求。

只用标准库（http.client），不跑真实模型回合，不联网连真实
GitHub/Slack（webhook 测试用的 secret 均为单测内临时值）。
每个测试用独立的临时 SQLite 文件与临时端口，跑完关闭 server。
"""

import hashlib
import hmac
import http.client
import json
import os
import sqlite3
import sys
import tempfile
import threading
import time
import unittest
from pathlib import Path
from unittest import mock

try:
    from cloud.auth import AuthStore
    from cloud.server import create_server, sweep_stale_tasks
except ImportError:  # discover -s cloud/tests 时把项目根目录补进 sys.path
    sys.path.insert(0, str(Path(__file__).resolve().parents[2]))
    from cloud.auth import AuthStore
    from cloud.server import create_server, sweep_stale_tasks


class ServerTestCase(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._old_admins = os.environ.get("CODEX_GLOBAL_ADMINS")
        os.environ["CODEX_GLOBAL_ADMINS"] = self.owner_id
        self.viewer_id, self.viewer_key = self.auth.create_user("viewer")
        self.outsider_id, self.outsider_key = self.auth.create_user("outsider")
        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()
        # owner 建一个项目，viewer 设为该项目的 viewer 角色。
        status, self.project = self._req(
            "POST", "/projects", {"name": "demo"}, self.owner_key)
        self.assertEqual(status, 201)
        self.auth.set_member(
            self.project["id"], self.owner_id, self.viewer_id, "viewer")

    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()
        if getattr(self, "_old_admins", None) is None:
            os.environ.pop("CODEX_GLOBAL_ADMINS", None)
        else:
            os.environ["CODEX_GLOBAL_ADMINS"] = self._old_admins

    # -- HTTP 工具 -----------------------------------------------------
    def _req(self, method, path, body=None, key=None, headers=None,
             raw_body=None):
        conn = http.client.HTTPConnection(
            "127.0.0.1", self.port, timeout=10)
        hdrs = dict(headers or {})
        if raw_body is not None:
            payload = raw_body
        elif body is not None:
            payload = json.dumps(body).encode("utf-8")
            hdrs.setdefault("Content-Type", "application/json")
        else:
            payload = None
        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 _create_task(self, key, **extra):
        body = {"project_id": self.project["id"], "prompt": "做一个加法函数"}
        body.update(extra)
        status, task = self._req("POST", "/tasks", body, key)
        self.assertEqual(status, 201, task)
        return task

    # -- 健康检查与认证 -------------------------------------------------
    def test_health_no_auth(self):
        status, body = self._req("GET", "/health")
        self.assertEqual(status, 200)
        # health 除 ok/service 外还返回 backend 与 db_ok（加法字段，
        # 不是改成功语义）；旧的精确全等断言已随字段扩展而放宽。
        self.assertTrue(body["ok"])
        self.assertEqual(body["service"], "cloud-control-plane")
        self.assertEqual(body["backend"], "sqlite")
        self.assertTrue(body["db_ok"])

    def test_no_key_401(self):
        status, body = self._req("GET", "/projects")
        self.assertEqual(status, 401)
        self.assertEqual(body, {"error": "unauthorized"})

    def test_wrong_key_401(self):
        status, body = self._req("GET", "/projects", key="wrong-key")
        self.assertEqual(status, 401)
        self.assertEqual(body, {"error": "unauthorized"})

    # -- 项目可见性与权限 -----------------------------------------------
    def test_projects_only_member_projects(self):
        # outsider 建自己的项目；双方互不可见对方的项目。
        status, other = self._req(
            "POST", "/projects", {"name": "other"}, self.outsider_key)
        self.assertEqual(status, 201)
        status, body = self._req("GET", "/projects", key=self.owner_key)
        self.assertEqual(status, 200)
        ids = [p["id"] for p in body["projects"]]
        self.assertIn(self.project["id"], ids)
        self.assertNotIn(other["id"], ids)
        status, body = self._req("GET", "/projects", key=self.outsider_key)
        self.assertEqual(status, 200)
        ids = [p["id"] for p in body["projects"]]
        self.assertIn(other["id"], ids)
        self.assertNotIn(self.project["id"], ids)

    def test_nonmember_read_project_tasks_403(self):
        status, body = self._req(
            "GET", "/tasks?project_id=%s" % self.project["id"],
            key=self.outsider_key)
        self.assertEqual(status, 403)

    def test_viewer_create_task_403(self):
        status, body = self._req(
            "POST", "/tasks",
            {"project_id": self.project["id"], "prompt": "x"},
            self.viewer_key)
        self.assertEqual(status, 403)

    def test_viewer_can_read_tasks(self):
        self._create_task(self.owner_key)
        status, body = self._req(
            "GET", "/tasks?project_id=%s" % self.project["id"],
            key=self.viewer_key)
        self.assertEqual(status, 200)
        self.assertEqual(len(body["tasks"]), 1)

    # -- 任务 CRUD、幂等、取消、心跳、回收 -------------------------------
    def test_create_and_get_task(self):
        task = self._create_task(self.owner_key, priority=3)
        self.assertEqual(task["status"], "queued")
        self.assertEqual(task["priority"], 3)
        status, got = self._req(
            "GET", "/tasks/%s" % task["id"], key=self.owner_key)
        self.assertEqual(status, 200)
        self.assertEqual(got["id"], task["id"])
        status, _ = self._req("GET", "/tasks/nope", key=self.owner_key)
        self.assertEqual(status, 404)
        status, _ = self._req(
            "GET", "/tasks/%s" % task["id"], key=self.outsider_key)
        self.assertEqual(status, 403)

    def test_idempotent_key_returns_same_task(self):
        first = self._create_task(
            self.owner_key, idempotency_key="key-1")
        status, second = self._req(
            "POST", "/tasks",
            {"project_id": self.project["id"], "prompt": "另一个 prompt",
             "idempotency_key": "key-1"},
            self.owner_key)
        self.assertEqual(status, 200)
        self.assertEqual(second["id"], first["id"])
        self.assertEqual(second["prompt"], first["prompt"])
        status, body = self._req(
            "GET", "/tasks?project_id=%s" % self.project["id"],
            key=self.owner_key)
        self.assertEqual(len(body["tasks"]), 1)

    def test_cancel_running_then_terminal_cancel_400(self):
        task = self._create_task(self.owner_key)
        status, running = self._req(
            "POST", "/tasks/%s/heartbeat" % task["id"], {}, self.owner_key)
        self.assertEqual(status, 200)
        self.assertEqual(running["status"], "running")
        status, cancelled = self._req(
            "POST", "/tasks/%s/cancel" % task["id"], {}, self.owner_key)
        self.assertEqual(status, 200)
        self.assertEqual(cancelled["status"], "interrupted")
        status, body = self._req(
            "POST", "/tasks/%s/cancel" % task["id"], {}, self.owner_key)
        self.assertEqual(status, 400)

    def test_cancel_queued_task(self):
        task = self._create_task(self.owner_key)
        status, body = self._req(
            "POST", "/tasks/%s/cancel" % task["id"], {}, self.owner_key)
        self.assertEqual(status, 200)
        self.assertEqual(body["status"], "interrupted")

    def test_heartbeat_after_sweep_keeps_fresh_task(self):
        task = self._create_task(self.owner_key)
        self._req("POST", "/tasks/%s/heartbeat" % task["id"], {},
                  self.owner_key)
        status, body = self._req(
            "POST", "/admin/sweep", {}, self.owner_key)
        self.assertEqual(status, 200)
        self.assertNotIn(task["id"], body["swept"])
        status, got = self._req(
            "GET", "/tasks/%s" % task["id"], key=self.owner_key)
        self.assertEqual(got["status"], "running")

    def test_sweep_stale_with_threshold_zero(self):
        task = self._create_task(self.owner_key)
        self._req("POST", "/tasks/%s/heartbeat" % task["id"], {},
                  self.owner_key)
        swept = sweep_stale_tasks(self.db_path, threshold_seconds=0)
        self.assertIn(task["id"], swept)
        status, got = self._req(
            "GET", "/tasks/%s" % task["id"], key=self.owner_key)
        self.assertEqual(got["status"], "interrupted")

    def test_sweep_stale_manually_aged_heartbeat(self):
        task = self._create_task(self.owner_key)
        self._req("POST", "/tasks/%s/heartbeat" % task["id"], {},
                  self.owner_key)
        conn = sqlite3.connect(self.db_path)
        try:
            conn.execute(
                "UPDATE tasks SET last_heartbeat = ? WHERE id = ?",
                (time.time() - 3600, task["id"]))
            conn.commit()
        finally:
            conn.close()
        status, body = self._req(
            "POST", "/admin/sweep", {}, self.owner_key)
        self.assertEqual(status, 200)
        self.assertIn(task["id"], body["swept"])
        status, got = self._req(
            "GET", "/tasks/%s" % task["id"], key=self.owner_key)
        self.assertEqual(got["status"], "interrupted")

    def test_sweep_skips_queued_task(self):
        task = self._create_task(self.owner_key)  # 仍是 queued
        swept = sweep_stale_tasks(self.db_path, threshold_seconds=0)
        self.assertNotIn(task["id"], swept)

    # -- webhooks（验签，解析但不执行）----------------------------------
    def test_github_webhook_secret_missing_503(self):
        with mock.patch.dict(os.environ, {}, clear=False):
            os.environ.pop("GITHUB_WEBHOOK_SECRET", None)
            status, _ = self._req(
                "POST", "/webhooks/github", {"action": "created"},
                key=None)
        self.assertEqual(status, 503)

    def test_github_webhook_bad_signature_401(self):
        payload = {"action": "created", "comment": {"body": "/fix-review"}}
        raw = json.dumps(payload).encode("utf-8")
        with mock.patch.dict(
                os.environ, {"GITHUB_WEBHOOK_SECRET": "test-secret"}):
            status, body = self._req(
                "POST", "/webhooks/github", raw_body=raw,
                headers={"Content-Type": "application/json",
                         "X-Hub-Signature-256": "sha256=deadbeef"})
        self.assertEqual(status, 401)

    def test_github_webhook_good_signature_parses(self):
        payload = {
            "action": "created",
            "comment": {"body": "请 /fix-review 处理", "user": {"login": "alice"}},
            "repository": {"full_name": "org/repo"},
            "issue": {"number": 7},
        }
        raw = json.dumps(payload).encode("utf-8")
        sig = "sha256=" + hmac.new(
            b"test-secret", raw, hashlib.sha256).hexdigest()
        with mock.patch.dict(
                os.environ, {"GITHUB_WEBHOOK_SECRET": "test-secret"}):
            status, body = self._req(
                "POST", "/webhooks/github", raw_body=raw,
                headers={"Content-Type": "application/json",
                         "X-Hub-Signature-256": sig})
        self.assertEqual(status, 200)
        self.assertEqual(body["type"], "fix_review")
        self.assertEqual(body["pr_number"], 7)

    def test_slack_webhook_bad_and_good_signature(self):
        raw = b"text=/codex status my-task"
        ts = str(int(time.time()))
        good_sig = "v0=" + hmac.new(
            b"slack-secret",
            b"v0:" + ts.encode() + b":" + raw, hashlib.sha256).hexdigest()
        headers = {"Content-Type": "application/x-www-form-urlencoded",
                   "X-Slack-Signature": "v0=deadbeef",
                   "X-Slack-Request-Timestamp": ts}
        with mock.patch.dict(
                os.environ, {"SLACK_SIGNING_SECRET": "slack-secret"}):
            status, _ = self._req(
                "POST", "/webhooks/slack", raw_body=raw, headers=headers)
            self.assertEqual(status, 401)
            headers["X-Slack-Signature"] = good_sig
            status, body = self._req(
                "POST", "/webhooks/slack", raw_body=raw, headers=headers)
        self.assertEqual(status, 200)
        self.assertEqual(body, {"command": "status", "task": "my-task"})


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