"""多用户生产级缺口批单测：sweep 鉴权、默认执行器、新用户默认配额、 无价模型 token 上限（服务端准入）。 只用本地临时端口 + 临时 SQLite，不跑真实模型，不联网。 环境变量（默认执行器、默认配额额、token 上限）是服务端每次 请求现读的，测试里用 try/finally 改回，不影响其他测试。 """ import http.client import json import os import sqlite3 import sys import tempfile import threading import time import unittest from datetime import datetime, timezone from pathlib import Path try: from cloud.auth import AuthStore from cloud.cost import CostStore from cloud.server import _apply_default_quota, _default_quota_usd_micros 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.cost import CostStore from cloud.server import _apply_default_quota, _default_quota_usd_micros from cloud.server import create_server def _req(port, method, path, body=None, key=None): conn = http.client.HTTPConnection("127.0.0.1", 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 _utc_today(): return datetime.now(timezone.utc).strftime("%Y-%m-%d") def _setenv(name, value): old = os.environ.get(name) if value is None: os.environ.pop(name, None) else: os.environ[name] = value return old class MultiuserGapsTestCase(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() s, proj = _req(self.port, "POST", "/projects", {"name": "gproj"}, self.owner_key) self.assertEqual(s, 201, proj) self.project_id = proj["id"] 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 # -- helpers ------------------------------------------------------- def _create_task(self, key=None, body_extra=None): body = {"project_id": self.project_id, "prompt": "gap probe"} if body_extra: body.update(body_extra) return _req(self.port, "POST", "/tasks", body, key if key is not None else self.owner_key) def _insert_quota(self, user_id, micros): conn = sqlite3.connect(self.db_path) try: conn.execute( "CREATE TABLE IF NOT EXISTS user_quotas(" "user_id TEXT PRIMARY KEY, quota_usd_micros BIGINT," " updated_at DOUBLE PRECISION)") conn.execute( "INSERT OR REPLACE INTO user_quotas" "(user_id, quota_usd_micros, updated_at) VALUES(?, ?, ?)", (user_id, micros, time.time())) conn.commit() finally: conn.close() def _insert_unpriced_cost(self, task_id, input_tokens, output_tokens): conn = sqlite3.connect(self.db_path) try: conn.execute( "INSERT INTO tasks(id, project_id, prompt, status," " idempotency_key, depends_on, priority, created_by," " created_at, updated_at, last_heartbeat)" " VALUES(?, ?, ?, 'completed', NULL, '[]', 0, ?, ?, ?, NULL)", (task_id, self.project_id, "donor", self.owner_id, time.time(), time.time())) conn.commit() finally: conn.close() # CostStore.record 建表并按 UTC 今天记 day，与服务端查询口径一致。 store = CostStore(self.db_path) try: store.record(task_id, self.project_id, "gpt-5.6-terra", input_tokens=input_tokens, output_tokens=output_tokens, usd_micros=None, source="actual") finally: try: store._conn.close() except Exception: pass def _make_running_task(self): s, task = self._create_task() self.assertEqual(s, 201, task) s, _ = _req(self.port, "POST", "/tasks/%s/heartbeat" % task["id"], {}, self.owner_key) self.assertEqual(s, 200) return task["id"] # -- sweep 鉴权 ----------------------------------------------------- def test_sweep_viewer_and_outsider_forbidden(self): for key in (self.viewer_key, self.outsider_key): s, body = _req(self.port, "POST", "/admin/sweep", {}, key) self.assertEqual(s, 403, body) def test_sweep_owner_maintainer_ok_and_sweeps_stale(self): tid = self._make_running_task() conn = sqlite3.connect(self.db_path) try: conn.execute( "UPDATE tasks SET last_heartbeat = ? WHERE id = ?", (time.time() - 3600, tid)) conn.commit() finally: conn.close() s, body = _req(self.port, "POST", "/admin/sweep", {"threshold_seconds": 300}, self.owner_key) self.assertEqual(s, 200, body) self.assertIn(tid, body["swept"]) def test_sweep_threshold_clamped_to_60(self): tid = self._make_running_task() s, body = _req(self.port, "POST", "/admin/sweep", {"threshold_seconds": 0}, self.owner_key) self.assertEqual(s, 200, body) self.assertNotIn(tid, body["swept"]) # -- 默认执行器 ------------------------------------------------------ def test_default_executor_is_docker(self): s, task = self._create_task() self.assertEqual(s, 201, task) self.assertEqual(task["executor"], "docker") def test_explicit_executor_kept(self): s, task = self._create_task(body_extra={"executor": "local"}) self.assertEqual(s, 201, task) self.assertEqual(task["executor"], "local") def test_default_executor_env_override(self): old = _setenv("CLOUD_DEFAULT_EXECUTOR", "local") try: s, task = self._create_task() self.assertEqual(s, 201, task) self.assertEqual(task["executor"], "local") finally: _setenv("CLOUD_DEFAULT_EXECUTOR", old) old = _setenv("CLOUD_DEFAULT_EXECUTOR", "") try: s, task = self._create_task() self.assertEqual(s, 201, task) self.assertIsNone(task["executor"]) finally: _setenv("CLOUD_DEFAULT_EXECUTOR", old) # -- 新用户默认配额 --------------------------------------------------- def test_apply_default_quota_idempotent(self): conn = sqlite3.connect(self.db_path) try: self.assertTrue(_apply_default_quota(conn, "newbie")) row = conn.execute( "SELECT quota_usd_micros FROM user_quotas" " WHERE user_id = 'newbie'").fetchone() self.assertEqual(row[0], _default_quota_usd_micros()) # 第二次不再写，不覆盖管理员后来改过的额度。 conn.execute( "UPDATE user_quotas SET quota_usd_micros = 12345" " WHERE user_id = 'newbie'") conn.commit() self.assertFalse(_apply_default_quota(conn, "newbie")) row = conn.execute( "SELECT quota_usd_micros FROM user_quotas" " WHERE user_id = 'newbie'").fetchone() self.assertEqual(row[0], 12345) finally: conn.close() def test_default_quota_env_override(self): old = _setenv("CLOUD_DEFAULT_QUOTA_USD_MICROS", "5000") try: self.assertEqual(_default_quota_usd_micros(), 5000) finally: _setenv("CLOUD_DEFAULT_QUOTA_USD_MICROS", old) # -- 无价 token 上限（服务端准入）-------------------------------------- def test_unpriced_token_cap_blocks_create(self): self._insert_quota(self.owner_id, 10 ** 12) # USD 很大，不挡 for i in range(7): # 7 * 150 = 1050 token self._insert_unpriced_cost("donor-%d" % i, 100, 50) old = _setenv("CLOUD_QUOTA_UNPRICED_TOKEN_CAP", "1000") try: s, body = self._create_task() self.assertEqual(s, 402, body) self.assertEqual(body["error"], "quota_exceeded") self.assertEqual(body["unpriced_tokens_used"], 1050) finally: _setenv("CLOUD_QUOTA_UNPRICED_TOKEN_CAP", old) def test_unpriced_token_cap_under_cap_passes(self): self._insert_quota(self.owner_id, 10 ** 12) for i in range(6): # 6 * 150 = 900 < 1000 self._insert_unpriced_cost("donor-%d" % i, 100, 50) old = _setenv("CLOUD_QUOTA_UNPRICED_TOKEN_CAP", "1000") try: s, task = self._create_task() self.assertEqual(s, 201, task) finally: _setenv("CLOUD_QUOTA_UNPRICED_TOKEN_CAP", old) def test_unpriced_token_cap_blocks_without_quota_row(self): # P0 安全批后：无配额行不再等于不限额，无价 cap 照样挡。 for i in range(7): self._insert_unpriced_cost("donor-%d" % i, 100, 50) old = _setenv("CLOUD_QUOTA_UNPRICED_TOKEN_CAP", "1000") try: s, body = self._create_task() self.assertEqual(s, 402, body) self.assertEqual(body["error"], "quota_exceeded") finally: _setenv("CLOUD_QUOTA_UNPRICED_TOKEN_CAP", old) if __name__ == "__main__": unittest.main()