多用户生产级缺口批：sweep 鉴权、默认配额、默认 executor、 无价 token 上限
- POST /admin/sweep 改走任一项目 maintainer+ 门槛（与备份管理同口径），否则 403，记 admin_sweep 审计；threshold 下限钳 60 秒 - OAuth 建号 created=True 时写默认配额行（默认 $10/月，CLOUD_DEFAULT_QUOTA_USD_MICROS 可调，INSERT OR IGNORE 幂等）；老用户无配额行不受影响 - 建任务未指定且无环境继承时默认 executor=docker（分支继承为空时同理）；CLOUD_DEFAULT_EXECUTOR 可覆盖，空字符串恢复旧行为；命令台下拉默认改为 docker - 无价模型 token 上限：价目表外模型（默认 terra）成本 usd 为 NULL，USD 配额永不生效；有配额行的用户当月无价 token 达 CLOUD_QUOTA_UNPRICED_TOKEN_CAP（默认 2000 万/月）即停，服务端 402、worker 认领/重试/后续消息同口径；报表仍记 unpriced，不折假美元 - 测试：新增 test_multiuser_gaps 11 个 + test_worker_quota 3 个；全量 cloud 668、host 64、deploy 19 全过
This commit is contained in:
host committed
1 parent
c1f9313c10
commit
7e081211db
7 files changed
+580
-13
No files matched your search
@@ -0,0 +1,279 @@
|
||||
"""多用户生产级缺口批单测：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.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()
|
||||
|
||||
# -- 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_tokens_unlimited_without_quota_row(self):
|
||||
for i in range(7):
|
||||
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)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
Reference in new issue
Block a user