租户配额与计费：按用户月度 USD 上限 + 建任务/分支/重试准入 402 + 管理/自查端点 + 控制台展示
This commit is contained in:
host committed
1 parent
2a7cdd582f
commit
970b181bab
5 files changed
+496
-1
No files matched your search
@@ -0,0 +1,239 @@
|
||||
"""租户配额与计费单测：月度 USD 上限 + 准入 402 + 管理端点。
|
||||
|
||||
不跑真实模型回合、不联网：server 侧走本地临时端口 + 临时
|
||||
SQLite；用量直接插 cost_records 行（按任务 created_by 归属、
|
||||
UTC day），验证 JOIN 归属、自然月窗口与三路准入（建任务 /
|
||||
分支 / 重试）。配额单位是微美元（1 美元 = 1_000_000），
|
||||
与 cost_records 一致。
|
||||
"""
|
||||
|
||||
import http.client
|
||||
import json
|
||||
import os
|
||||
import sqlite3
|
||||
import sys
|
||||
import tempfile
|
||||
import threading
|
||||
import unittest
|
||||
from datetime import date, datetime, timedelta, timezone
|
||||
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
|
||||
|
||||
|
||||
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 _set_task_status(db_path, task_id, status):
|
||||
conn = sqlite3.connect(db_path)
|
||||
try:
|
||||
conn.execute("UPDATE tasks SET status = ? WHERE id = ?",
|
||||
(status, task_id))
|
||||
conn.commit()
|
||||
finally:
|
||||
conn.close()
|
||||
|
||||
|
||||
def _insert_cost(db_path, task_id, project_id, usd_micros, day):
|
||||
conn = sqlite3.connect(db_path)
|
||||
try:
|
||||
conn.execute(
|
||||
"INSERT INTO cost_records(task_id, project_id, model,"
|
||||
" input_tokens, output_tokens, total_tokens, usd_micros,"
|
||||
" source, day, created_at) VALUES(?, ?, ?, ?, ?, ?, ?, ?, ?, ?)",
|
||||
(task_id, project_id, "gpt-5.6-terra", 100, 50, 150,
|
||||
usd_micros, "actual", day, 0))
|
||||
conn.commit()
|
||||
finally:
|
||||
conn.close()
|
||||
|
||||
|
||||
def _utc_today():
|
||||
return datetime.now(timezone.utc).strftime("%Y-%m-%d")
|
||||
|
||||
|
||||
def _prev_month_day():
|
||||
today = date.today()
|
||||
first = today.replace(day=1)
|
||||
return (first - timedelta(days=1)).isoformat()
|
||||
|
||||
|
||||
class ServerQuotasTestCase(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.other_id, self.other_key = self.auth.create_user("other")
|
||||
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": "qproj"}, self.owner_key)
|
||||
self.assertEqual(s, 201, proj)
|
||||
self.project_id = proj["id"]
|
||||
|
||||
def tearDown(self):
|
||||
self.server.shutdown()
|
||||
self.server.server_close()
|
||||
self.tmp.cleanup()
|
||||
|
||||
def _create_task(self, key=None):
|
||||
return _req(self.port, "POST", "/tasks",
|
||||
{"project_id": self.project_id,
|
||||
"prompt": "quota probe"},
|
||||
key if key is not None else self.owner_key)
|
||||
|
||||
def _set_quota(self, target, quota, key=None):
|
||||
return _req(self.port, "POST", "/admin/quotas",
|
||||
{"user_id": target, "quota_usd_micros": quota},
|
||||
key if key is not None else self.owner_key)
|
||||
|
||||
def test_self_quota_default_unlimited(self):
|
||||
s, q = _req(self.port, "GET", "/users/me/quota",
|
||||
key=self.owner_key)
|
||||
self.assertEqual(s, 200, q)
|
||||
self.assertIsNone(q["quota_usd_micros"])
|
||||
self.assertEqual(q["used_usd_micros"], 0)
|
||||
self.assertIsNone(q["remaining_usd_micros"])
|
||||
self.assertEqual(q["month"],
|
||||
datetime.now(timezone.utc).strftime("%Y-%m"))
|
||||
|
||||
def test_admin_set_and_get(self):
|
||||
s, q = self._set_quota(self.owner_id, 5000000)
|
||||
self.assertEqual(s, 200, q)
|
||||
self.assertEqual(q["quota_usd_micros"], 5000000)
|
||||
s, q = _req(self.port, "GET",
|
||||
"/admin/quotas?user_id=" + self.owner_id,
|
||||
key=self.owner_key)
|
||||
self.assertEqual(s, 200, q)
|
||||
self.assertEqual(q["quota_usd_micros"], 5000000)
|
||||
s, q = _req(self.port, "GET", "/users/me/quota",
|
||||
key=self.owner_key)
|
||||
self.assertEqual(q["quota_usd_micros"], 5000000)
|
||||
# null = 恢复不限额
|
||||
s, q = self._set_quota(self.owner_id, None)
|
||||
self.assertEqual(s, 200, q)
|
||||
self.assertIsNone(q["quota_usd_micros"])
|
||||
|
||||
def test_admin_set_forbidden_and_validation(self):
|
||||
# other 没有任何项目角色，不是管理员
|
||||
s, q = self._set_quota(self.owner_id, 100, self.other_key)
|
||||
self.assertEqual(s, 403, q)
|
||||
s, q = _req(self.port, "GET",
|
||||
"/admin/quotas?user_id=" + self.owner_id,
|
||||
key=self.other_key)
|
||||
self.assertEqual(s, 403, q)
|
||||
s, q = self._set_quota(self.owner_id, -1)
|
||||
self.assertEqual(s, 400, q)
|
||||
s, q = self._set_quota("no-such-user", 100)
|
||||
self.assertEqual(s, 404, q)
|
||||
s, q = _req(self.port, "POST", "/admin/quotas",
|
||||
{"quota_usd_micros": 100}, self.owner_key)
|
||||
self.assertEqual(s, 400, q)
|
||||
s, q = _req(self.port, "GET", "/admin/quotas",
|
||||
key=self.owner_key)
|
||||
self.assertEqual(s, 400, q)
|
||||
|
||||
def test_usage_counts_current_month_only(self):
|
||||
s, task = self._create_task()
|
||||
self.assertEqual(s, 201, task)
|
||||
_insert_cost(self.db_path, task["id"], self.project_id,
|
||||
5000000, _utc_today())
|
||||
_insert_cost(self.db_path, task["id"], self.project_id,
|
||||
100000000, _prev_month_day())
|
||||
s, q = _req(self.port, "GET", "/users/me/quota",
|
||||
key=self.owner_key)
|
||||
self.assertEqual(s, 200, q)
|
||||
self.assertEqual(q["used_usd_micros"], 5000000)
|
||||
|
||||
def test_create_blocked_at_quota(self):
|
||||
s, task = self._create_task()
|
||||
self.assertEqual(s, 201, task)
|
||||
_insert_cost(self.db_path, task["id"], self.project_id,
|
||||
5000000, _utc_today())
|
||||
s, _ = self._set_quota(self.owner_id, 5000000)
|
||||
self.assertEqual(s, 200)
|
||||
s, body = self._create_task()
|
||||
self.assertEqual(s, 402, body)
|
||||
self.assertEqual(body["error"], "quota_exceeded")
|
||||
self.assertEqual(body["quota_usd_micros"], 5000000)
|
||||
self.assertEqual(body["used_usd_micros"], 5000000)
|
||||
# 额度调高后放行
|
||||
s, _ = self._set_quota(self.owner_id, 6000000)
|
||||
self.assertEqual(s, 200)
|
||||
s, task2 = self._create_task()
|
||||
self.assertEqual(s, 201, task2)
|
||||
|
||||
def test_fork_blocked_at_quota(self):
|
||||
s, task = self._create_task()
|
||||
self.assertEqual(s, 201, task)
|
||||
_insert_cost(self.db_path, task["id"], self.project_id,
|
||||
3000000, _utc_today())
|
||||
_set_task_status(self.db_path, task["id"], "completed")
|
||||
s, _ = self._set_quota(self.owner_id, 3000000)
|
||||
self.assertEqual(s, 200)
|
||||
s, body = _req(self.port, "POST",
|
||||
"/tasks/%s/fork" % task["id"],
|
||||
{"message": "branch probe"}, self.owner_key)
|
||||
self.assertEqual(s, 402, body)
|
||||
s, _ = self._set_quota(self.owner_id, None)
|
||||
self.assertEqual(s, 200)
|
||||
s, forked = _req(self.port, "POST",
|
||||
"/tasks/%s/fork" % task["id"],
|
||||
{"message": "branch probe"}, self.owner_key)
|
||||
self.assertEqual(s, 201, forked)
|
||||
|
||||
def test_retry_blocked_at_quota(self):
|
||||
s, task = self._create_task()
|
||||
self.assertEqual(s, 201, task)
|
||||
_insert_cost(self.db_path, task["id"], self.project_id,
|
||||
2000000, _utc_today())
|
||||
_set_task_status(self.db_path, task["id"], "failed")
|
||||
s, _ = self._set_quota(self.owner_id, 2000000)
|
||||
self.assertEqual(s, 200)
|
||||
s, body = _req(self.port, "POST",
|
||||
"/tasks/%s/retry" % task["id"], {},
|
||||
self.owner_key)
|
||||
self.assertEqual(s, 402, body)
|
||||
s, row = _req(self.port, "GET", "/tasks/" + task["id"],
|
||||
key=self.owner_key)
|
||||
self.assertEqual(row["status"], "failed")
|
||||
s, _ = self._set_quota(self.owner_id, None)
|
||||
self.assertEqual(s, 200)
|
||||
s, retried = _req(self.port, "POST",
|
||||
"/tasks/%s/retry" % task["id"], {},
|
||||
self.owner_key)
|
||||
self.assertEqual(s, 200, retried)
|
||||
self.assertEqual(retried["status"], "queued")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
Reference in new issue
Block a user