240 lines
9.1 KiB
Python
240 lines
9.1 KiB
Python
"""租户配额与计费单测：月度 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()
|