126 lines
4.5 KiB
Python
126 lines
4.5 KiB
Python
"""并发上限服务端面单测（/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")
|
||
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 _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)
|
||
# 读权限：任何已认证用户可读（viewer 也能看到基础设施行）。
|
||
workers = self._workers(self.viewer_key)
|
||
self.assertEqual(workers["w1"]["capacity"], 4)
|
||
self.assertEqual(workers["w1"]["running"], 2)
|
||
|
||
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()
|