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

"""并发上限服务端面单测（/workers 占用展示）。

覆盖：GET /workers 行里含 capacity 与 running；心跳上报 running
后列表可见；心跳未给 running 时保留旧值；非法 running 400；
未认证 401；/workers 读权限口径是任何已认证用户（与既有
_workers_list 一致，不按项目成员过滤）。

只用标准库 http.client，不跑真实模型，不联网。
"""

import http.client
import json
import os
import sys
import tempfile
import threading
import unittest
from pathlib import Path

try:
    from cloud.auth import AuthStore
    from cloud.server import create_server
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


class CapsServerTestCase(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")
        # P1：完整 workers 清单只限全局管理员；owner 在此测试里即管理员。
        self._old_admins = os.environ.get("CODEX_GLOBAL_ADMINS")
        os.environ["CODEX_GLOBAL_ADMINS"] = self.owner_id
        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()
        if getattr(self, "_old_admins", None) is None:
            os.environ.pop("CODEX_GLOBAL_ADMINS", None)
        else:
            os.environ["CODEX_GLOBAL_ADMINS"] = self._old_admins

    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 _workers(self, key):
        status, body = self._req("GET", "/workers", key=key)
        self.assertEqual(status, 200, body)
        return {w["worker_id"]: w for w in body["workers"]}

    def test_workers_list_shows_capacity_and_running(self):
        status, body = self._req(
            "POST", "/workers/heartbeat",
            {"worker_id": "w1", "executor": "local",
             "capacity": 4, "running": 2},
            self.owner_key)
        self.assertEqual(status, 200, body)
        # P1：完整清单只限全局管理员（owner 在此测试里即管理员）。
        workers = self._workers(self.owner_key)
        self.assertEqual(workers["w1"]["capacity"], 4)
        self.assertEqual(workers["w1"]["running"], 2)

    def test_workers_list_non_admin_gets_aggregate_only(self):
        self._req("POST", "/workers/heartbeat",
                  {"worker_id": "w1", "executor": "local",
                   "capacity": 4, "running": 2},
                  self.owner_key)
        status, body = self._req("GET", "/workers", key=self.viewer_key)
        self.assertEqual(status, 200, body)
        self.assertNotIn("workers", body)
        self.assertIn("workers_total", body)

    def test_heartbeat_without_running_preserves_old(self):
        self._req("POST", "/workers/heartbeat",
                  {"worker_id": "w1", "capacity": 2, "running": 1},
                  self.owner_key)
        status, body = self._req(
            "POST", "/workers/heartbeat", {"worker_id": "w1"},
            self.owner_key)
        self.assertEqual(status, 200, body)
        self.assertEqual(body["running"], 1)
        self.assertEqual(self._workers(self.owner_key)["w1"]["running"], 1)

    def test_heartbeat_invalid_running_400(self):
        status, _ = self._req(
            "POST", "/workers/heartbeat",
            {"worker_id": "w1", "running": "abc"}, self.owner_key)
        self.assertEqual(status, 400)
        status, _ = self._req(
            "POST", "/workers/heartbeat",
            {"worker_id": "w1", "running": -1}, self.owner_key)
        self.assertEqual(status, 400)

    def test_workers_requires_auth(self):
        status, _ = self._req("GET", "/workers")
        self.assertEqual(status, 401)

    def test_heartbeat_registers_unknown_worker_with_defaults(self):
        status, body = self._req(
            "POST", "/workers/heartbeat", {"worker_id": "new1"},
            self.owner_key)
        self.assertEqual(status, 200, body)
        self.assertEqual(body["capacity"], 1)
        self.assertEqual(body["running"], 0)


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