host
9f49b97cd3
- test_system_quota_exempt_server → test_system_quota_not_exempt_server：system 按普通属主算，quota 0 时照常拦（402）。 - test_system_task_exempt_from_quota → test_system_task_not_exempt_from_quota：system quota 0 时认领门拦，标 interrupted。 - test_poll_quota_lookup_failure_open → test_poll_quota_lookup_failure_open_returns_check_failed：open 异常时回 quota_check_failed（边界拒绝、轮询经 _quota_hit 过滤只警告），closed 回 quota_check_unavailable。
457 lines
19 KiB
Python
457 lines
19 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
|
||
from types import SimpleNamespace
|
||
from unittest import mock
|
||
|
||
try:
|
||
from cloud.auth import AuthStore
|
||
from cloud.server import CloudControlPlaneHandler, 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 CloudControlPlaneHandler, 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")
|
||
# P0 安全批后：admin 端点只认 CODEX_GLOBAL_ADMINS，
|
||
# 本文件里 owner 即管理员，测前写 env、测后还原。
|
||
self._old_admins = os.environ.get("CODEX_GLOBAL_ADMINS")
|
||
os.environ["CODEX_GLOBAL_ADMINS"] = self.owner_id
|
||
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()
|
||
if self._old_admins is None:
|
||
os.environ.pop("CODEX_GLOBAL_ADMINS", None)
|
||
else:
|
||
os.environ["CODEX_GLOBAL_ADMINS"] = self._old_admins
|
||
|
||
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_is_default_quota_not_unlimited(self):
|
||
# P0 安全批后：无配额行按默认额度处理，不再是 None。
|
||
s, q = _req(self.port, "GET", "/users/me/quota",
|
||
key=self.owner_key)
|
||
self.assertEqual(s, 200, q)
|
||
self.assertEqual(q["quota_usd_micros"], 10000000)
|
||
self.assertEqual(q["used_usd_micros"], 0)
|
||
self.assertEqual(q["remaining_usd_micros"], 10000000)
|
||
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")
|
||
|
||
def test_retry_blocked_when_operator_or_creator_over_quota(self):
|
||
# P0 安全批后：重试按双人门查配额（触发人 + 创建人）。
|
||
# #90 的“只查创建者、操作人超限不拦”已被 P0 有意收紧：
|
||
# 操作人用自己的额度过门、费用记创建者名下，仍是绕限，
|
||
# 所以操作人超限也必须拦。other 须是 maintainer 才能过
|
||
# approve_merge 权限门，否则 403 会掩盖配额判断。
|
||
self.auth.set_member(self.project_id, self.owner_id,
|
||
self.other_id, "maintainer")
|
||
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.other_key)
|
||
self.assertEqual(s, 402, body)
|
||
# 创建者不限额、操作人额度为 0：P0 双人门应拦（#90
|
||
# 只查创建者的旧口径会放行，那正是 P0 要堵的绕限）。
|
||
s, _ = self._set_quota(self.owner_id, None)
|
||
self.assertEqual(s, 200)
|
||
s, _ = self._set_quota(self.other_id, 0)
|
||
self.assertEqual(s, 200)
|
||
s, body = _req(self.port, "POST",
|
||
"/tasks/%s/retry" % task["id"], {},
|
||
self.other_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.other_id, None)
|
||
self.assertEqual(s, 200)
|
||
s, retried = _req(self.port, "POST",
|
||
"/tasks/%s/retry" % task["id"], {},
|
||
self.other_key)
|
||
self.assertEqual(s, 200, retried)
|
||
self.assertEqual(retried["status"], "queued")
|
||
|
||
def test_continue_blocked_at_creator_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())
|
||
_set_task_status(self.db_path, task["id"], "completed")
|
||
s, _ = self._set_quota(self.owner_id, 5000000)
|
||
self.assertEqual(s, 200)
|
||
# 以前：先回 200 排队、随后认领门悄悄标 interrupted。
|
||
# 现在：API 侧直接 402，状态保持 completed 不变。
|
||
s, body = _req(self.port, "POST",
|
||
"/tasks/%s/continue" % task["id"],
|
||
{"message": "again"}, 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"], "completed")
|
||
# 额度调高后放行。
|
||
s, _ = self._set_quota(self.owner_id, None)
|
||
self.assertEqual(s, 200)
|
||
s, continued = _req(self.port, "POST",
|
||
"/tasks/%s/continue" % task["id"],
|
||
{"message": "again"}, self.owner_key)
|
||
self.assertEqual(s, 200, continued)
|
||
self.assertEqual(continued["status"], "queued")
|
||
self.auth.set_member(self.project_id, self.owner_id,
|
||
self.other_id, "maintainer")
|
||
s, task2 = self._create_task()
|
||
self.assertEqual(s, 201, task2)
|
||
_set_task_status(self.db_path, task2["id"], "completed")
|
||
s, _ = self._set_quota(self.other_id, 0)
|
||
self.assertEqual(s, 200)
|
||
# P0 安全批后：续跑按双人门查，操作人额度为 0、
|
||
# 创建者不限额时也应拦（#90 只查创建者的旧口径会
|
||
# 放行，那正是 P0 要堵的绕限）；两人都不超限才放行。
|
||
s, body = _req(self.port, "POST",
|
||
"/tasks/%s/continue" % task2["id"],
|
||
{"message": "again"}, self.other_key)
|
||
self.assertEqual(s, 402, body)
|
||
s, row = _req(self.port, "GET", "/tasks/" + task2["id"],
|
||
key=self.owner_key)
|
||
self.assertEqual(row["status"], "completed")
|
||
s, _ = self._set_quota(self.other_id, None)
|
||
self.assertEqual(s, 200)
|
||
s, continued2 = _req(self.port, "POST",
|
||
"/tasks/%s/continue" % task2["id"],
|
||
{"message": "again"}, self.other_key)
|
||
self.assertEqual(s, 200, continued2)
|
||
self.assertEqual(continued2["status"], "queued")
|
||
|
||
def test_quota_reports_unpriced_tokens(self):
|
||
s, task = self._create_task()
|
||
self.assertEqual(s, 201, task)
|
||
# 有价 5000000 计入 used；无价 actual 150 只进
|
||
# unpriced_actual_tokens，不计价、不拦任务。
|
||
_insert_cost(self.db_path, task["id"], self.project_id,
|
||
5000000, _utc_today())
|
||
_insert_cost(self.db_path, task["id"], self.project_id,
|
||
None, _utc_today())
|
||
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)
|
||
self.assertEqual(q["unpriced_actual_tokens"], 150)
|
||
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["unpriced_actual_tokens"], 150)
|
||
|
||
# -- G7：准入点 fail-closed ----------------------------------------
|
||
def test_create_quota_lookup_failure_503(self):
|
||
with mock.patch.object(
|
||
CloudControlPlaneHandler, "_quota_status",
|
||
side_effect=RuntimeError("db down")):
|
||
s, body = self._create_task()
|
||
self.assertEqual(s, 503, body)
|
||
self.assertEqual(body["error"], "quota_check_unavailable")
|
||
# 没建出任务：状态、任务表都不改。
|
||
conn = sqlite3.connect(self.db_path)
|
||
try:
|
||
n = conn.execute(
|
||
"SELECT COUNT(*) FROM tasks").fetchone()[0]
|
||
finally:
|
||
conn.close()
|
||
self.assertEqual(n, 0)
|
||
|
||
# -- G13：system 豁免（server 侧） -----------------------------------
|
||
def test_system_quota_not_exempt_server(self):
|
||
# P1 口径：system 按普通属主算，不豁免（G14 的 system 豁免
|
||
# 与 P1 冲突，rebase 时未合入）。quota 0、used 0 时
|
||
# used >= quota，照常拦。
|
||
conn = sqlite3.connect(self.db_path)
|
||
try:
|
||
conn.execute(
|
||
"INSERT INTO user_quotas(user_id, quota_usd_micros,"
|
||
" updated_at) VALUES('system', 0, 0)")
|
||
conn.commit()
|
||
finally:
|
||
conn.close()
|
||
handler = object.__new__(CloudControlPlaneHandler)
|
||
handler.server = SimpleNamespace(db_path=self.db_path)
|
||
handler._send_json = lambda *args, **kwargs: None
|
||
self.assertIsNotNone(handler._check_quota("system"))
|
||
self.assertTrue(handler._reject_if_over_quota("system"))
|
||
|
||
# -- G13/G14：/pr 按请求人过配额门、队列记归属人 ----------------------
|
||
def _merged_remote_task(self):
|
||
s, task = self._create_task()
|
||
self.assertEqual(s, 201, task)
|
||
conn = sqlite3.connect(self.db_path)
|
||
try:
|
||
conn.execute(
|
||
"UPDATE tasks SET status = 'merged',"
|
||
" worktree_host = 'elsewhere-host-1' WHERE id = ?",
|
||
(task["id"],))
|
||
conn.commit()
|
||
finally:
|
||
conn.close()
|
||
return task
|
||
|
||
def test_pr_blocked_at_quota_no_job(self):
|
||
task = self._merged_remote_task()
|
||
s, _ = self._set_quota(self.owner_id, 0)
|
||
self.assertEqual(s, 200)
|
||
s, body = _req(self.port, "POST",
|
||
"/tasks/%s/pr" % task["id"], {},
|
||
self.owner_key)
|
||
self.assertEqual(s, 402, body)
|
||
self.assertEqual(body["error"], "quota_exceeded")
|
||
conn = sqlite3.connect(self.db_path)
|
||
try:
|
||
has_table = conn.execute(
|
||
"SELECT 1 FROM sqlite_master WHERE type='table'"
|
||
" AND name='pr_jobs'").fetchone() is not None
|
||
rows = conn.execute(
|
||
"SELECT COUNT(*) FROM pr_jobs"
|
||
" WHERE task_id = ?", (task["id"],)).fetchone()[0] \
|
||
if has_table else 0
|
||
status = conn.execute(
|
||
"SELECT status FROM tasks WHERE id = ?",
|
||
(task["id"],)).fetchone()[0]
|
||
finally:
|
||
conn.close()
|
||
self.assertEqual(rows, 0)
|
||
self.assertEqual(status, "merged")
|
||
|
||
def test_pr_queue_records_created_by(self):
|
||
task = self._merged_remote_task()
|
||
s, body = _req(self.port, "POST",
|
||
"/tasks/%s/pr" % task["id"], {},
|
||
self.owner_key)
|
||
self.assertEqual(s, 200, body)
|
||
self.assertTrue(body.get("pr_queued"))
|
||
conn = sqlite3.connect(self.db_path)
|
||
try:
|
||
row = conn.execute(
|
||
"SELECT created_by FROM pr_jobs WHERE task_id = ?",
|
||
(task["id"],)).fetchone()
|
||
finally:
|
||
conn.close()
|
||
self.assertIsNotNone(row)
|
||
self.assertEqual(row[0], self.owner_id)
|
||
|
||
|
||
if __name__ == "__main__":
|
||
unittest.main()
|