From 970b181baba33c4e9b2c5f787c6b4bb3a27e793d Mon Sep 17 00:00:00 2001 From: host Date: Tue, 27 Oct 2026 05:18:37 +0800 Subject: [PATCH] =?UTF-8?q?=E7=A7=9F=E6=88=B7=E9=85=8D=E9=A2=9D=E4=B8=8E?= =?UTF-8?q?=E8=AE=A1=E8=B4=B9=EF=BC=9A=E6=8C=89=E7=94=A8=E6=88=B7=E6=9C=88?= =?UTF-8?q?=E5=BA=A6=20USD=20=E4=B8=8A=E9=99=90=20+=20=E5=BB=BA=E4=BB=BB?= =?UTF-8?q?=E5=8A=A1/=E5=88=86=E6=94=AF/=E9=87=8D=E8=AF=95=E5=87=86?= =?UTF-8?q?=E5=85=A5=20402=20+=20=E7=AE=A1=E7=90=86/=E8=87=AA=E6=9F=A5?= =?UTF-8?q?=E7=AB=AF=E7=82=B9=20+=20=E6=8E=A7=E5=88=B6=E5=8F=B0=E5=B1=95?= =?UTF-8?q?=E7=A4=BA?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- cloud/README.md | 41 ++++++ cloud/console.html | 23 ++- cloud/db.py | 4 + cloud/server.py | 190 +++++++++++++++++++++++++ cloud/tests/test_task_quotas.py | 239 ++++++++++++++++++++++++++++++++ 5 files changed, 496 insertions(+), 1 deletion(-) create mode 100644 cloud/tests/test_task_quotas.py diff --git a/cloud/README.md b/cloud/README.md index 5f3f96e..caf0add 100644 --- a/cloud/README.md +++ b/cloud/README.md @@ -871,6 +871,47 @@ interrupted 且终态任务 cancel 返回 400、心跳后 sweep 不扫刚心跳 `/tasks`、`/projects`、`/workers`、`/health` 同款，否则页面能打开 但环境区永远加载失败）。 +## 租户配额与计费 + +- 口径：配额是按用户、按 UTC 自然月的 USD 上限（`user_quotas` + 表，`quota_usd_micros` 为 NULL 表示不限额，0 表示一分钱也不 + 能花；金额单位与 cost_records 一致，都是微美元，避免浮点）。 + 用量只算 `usd_micros` 非空的实际成本行（estimated 行 usd 为 + NULL，SUM 自动跳过），按任务创建者 `created_by` 归属——一个 + 用户在多个项目里花的钱都算进他自己的月度账单；项目维度的 + 拆分由项目成本接口另行提供。 +- 准入：建任务、会话分支、重试在权限校验通过后、写库之前查 + 配额，本月用量 >= 配额时回 402 `quota_exceeded`（附 month、 + quota_usd_micros、used_usd_micros），不建任务、不改状态。 + 老用户与未设配额的用户默认不限额，行为与过去一致。注意 + 边界：只做准入校验，不做跑中掐断——已在跑的回合、后续消 + 息触发的新一轮中途超限不会被停掉（worker 侧没有按金额中途 + 停转的钩子）；重试按操作人查配额、费用仍记到原任务创建者 + 名下（建任务与分支的操作人就是创建者，没有这个错位）。 + 系统任务（GitHub webhook 的 /fix-review，created_by + 是 'system'）不记到任何用户名下、不占配额；要把它算到项 + 目所有者头上需另做归属。 +- 接口：`GET /users/me/quota`（自查，返回 month、quota、 + used、remaining，不限额时 quota 与 remaining 为 null）； + `GET /admin/quotas?user_id=` 与 `POST /admin/quotas` + （body：user_id、quota_usd_micros，null=不限额，负数 400， + 用户不存在 404，改完记 `quota_set` 审计）。管理口径沿用 + `/admin/backup` 的做法：在任一项目里是 maintainer/owner 即 + 算管理员，没有任何项目角色的用户无权改任何人的配额。 + 这还不是真正的收款：没有接支付网关、不开发票、不扣款， + 只有记账 + 配额 + 超限拦截；要收钱需另接支付（价格表在 + `cloud/cost.py`，注意那是 OpenAI 牌价、不是中转站的账单）。 +- 命令台：API Key 卡下方显示本月用量与配额（10 秒自动刷新； + 超限建任务时接口的 402 会带“本月配额已用完”提示）。管理 + 员改配额目前只走 API，控制台没有改配额的表单。 +- 公网 nginx：注意，443 块的默认 `location /` 是代理到主机 + `127.0.0.1:28080`（不是控制面），而 `/users`、`/admin` + 目前没有专属 location——配额接口部署上生产时须同批加 + `location ^~ /users/me/quota` 与 `location ^~ /admin/quotas` + （只放这两个精确前缀，不要整段放行 `/users` 或 `/admin`， + 否则会把 `/admin/sweep`、`/admin/backup` 也暴露到公网）， + 否则公网页面调配额接口会落到 28080 而不是控制面。 + ## 运维批（本轮新增：部署、恢复演练、监控、心跳、多用户验证） 细则在 `deploy/README.md`，这里只记代码侧的接口与默认值： diff --git a/cloud/console.html b/cloud/console.html index 99db94c..e6c1988 100644 --- a/cloud/console.html +++ b/cloud/console.html @@ -87,6 +87,7 @@ +
未加载
@@ -292,6 +293,7 @@ async function api(path, method, body) { var msg = (data && data.error) ? data.error : ("HTTP " + r.status); if (r.status === 401) msg += "（key 未填或无效）"; if (r.status === 429) msg += "（触发限流，稍后再试）"; + if (r.status === 402) msg += "（本月配额已用完，联系管理员调整）"; throw new Error(msg); } return data; @@ -324,6 +326,23 @@ async function loadHealth() { } catch (e) { $("#health").textContent = "加载失败：" + e.message; } } +async function loadQuota() { + // 本月配额与用量：只读展示，管理员改配额走 API（/admin/quotas）。 + var el = $("#quotaState"); + if (!getKey()) { el.textContent = "未填 key，配额不显示"; return; } + try { + var d = await api("/users/me/quota", "GET"); + var used = (d.used_usd_micros / 1000000).toFixed(2); + if (d.quota_usd_micros === null || d.quota_usd_micros === undefined) { + el.textContent = "本月用量 $" + used + "（不限额）"; + } else { + var q = (d.quota_usd_micros / 1000000).toFixed(2); + var rem = (d.remaining_usd_micros / 1000000).toFixed(2); + el.textContent = "本月用量 $" + used + " / $" + q + "（剩余 $" + rem + "）"; + } + } catch (e) { el.textContent = "配额加载失败：" + e.message; } +} + async function loadProjects() { try { var d = await api("/projects", "GET"); @@ -610,6 +629,7 @@ $("#saveKey").onclick = function () { loadWorkers(); if ($("#projectSelect").value) { loadTasks(); loadEnvironments(); } if (getKey()) loadRelayModels(); + loadQuota(); }; $("#clearKey").onclick = function () { localStorage.removeItem(KEY_STORE); setKeyState(); }; $("#refreshHealth").onclick = loadHealth; @@ -789,11 +809,12 @@ setKeyState(); loadHealth(); loadProjects(); loadWorkers(); +loadQuota(); // 回访用户 key 已在 localStorage 里：初始化也要加载模型列表， // 否则下拉永远只有默认项（此前只在点“保存”时调一次， // 2026-10-27 实测：存过 key 后重开页面模型选不了）。 if (getKey()) loadRelayModels(); -setInterval(function () { if ($("#projectSelect").value) { loadTasks(); loadEnvironments(); } loadHealth(); }, 10000); +setInterval(function () { if ($("#projectSelect").value) { loadTasks(); loadEnvironments(); } loadHealth(); loadQuota(); }, 10000); diff --git a/cloud/db.py b/cloud/db.py index c7afcbe..b66d198 100644 --- a/cloud/db.py +++ b/cloud/db.py @@ -885,6 +885,10 @@ POSTGRES_SCHEMA = ( "CREATE UNIQUE INDEX IF NOT EXISTS idx_env_name" " ON environments(project_id, name)" " WHERE status != 'archived'", + "CREATE TABLE IF NOT EXISTS user_quotas(" + "user_id TEXT PRIMARY KEY, " + "quota_usd_micros BIGINT, " + "updated_at DOUBLE PRECISION)", "CREATE TABLE IF NOT EXISTS summaries(" "id BIGINT GENERATED BY DEFAULT AS IDENTITY PRIMARY KEY, " "project_id TEXT, " diff --git a/cloud/server.py b/cloud/server.py index 21a1a99..de042b5 100644 --- a/cloud/server.py +++ b/cloud/server.py @@ -623,6 +623,19 @@ def _ensure_environments_table(conn): " WHERE status != 'archived'") +def _ensure_quotas_table(conn): + """幂等建配额表：server init_db 与配额查询都可调它，老库 + 第一次跑到这里时自动补表，不丢数据。quota_usd_micros 为 + NULL 表示不限额（与 0 区分：0 是一分钱也不能花）；金额 + 单位与 cost_records 一致，都是 USD 微单位（1 美元 = + 1_000_000），避免浮点。""" + conn.execute( + "CREATE TABLE IF NOT EXISTS user_quotas(" + "user_id TEXT PRIMARY KEY, " + "quota_usd_micros BIGINT, " + "updated_at DOUBLE PRECISION)") + + _ENV_BRANCH_RE = re.compile(r"^[A-Za-z0-9_.\-/]+$") @@ -683,6 +696,7 @@ def init_db(db_path): migrate_projects_table(conn) _ensure_followups_table(conn) _ensure_environments_table(conn) + _ensure_quotas_table(conn) _ensure_uploads_tables(conn) conn.commit() return @@ -714,6 +728,7 @@ def init_db(db_path): migrate_projects_table(conn) _ensure_followups_table(conn) _ensure_environments_table(conn) + _ensure_quotas_table(conn) _ensure_uploads_tables(conn) # worker 每轮 poll 都按 status 取任务（见 load_tasks 的 # WHERE status IN (...)），无索引时是全表扫描；任务表随 @@ -1098,6 +1113,8 @@ class CloudControlPlaneHandler(BaseHTTPRequestHandler): self._archive_environment(user_id, rest[: -len("/archive")]) else: self._send_json(404, {"error": "not found"}) + elif path == "/users/me/quota" and method == "GET": + self._user_quota_self(user_id) elif path == "/tasks" and method == "POST": self._create_task(user_id) elif path == "/uploads" and method == "POST": @@ -1108,6 +1125,10 @@ class CloudControlPlaneHandler(BaseHTTPRequestHandler): self._admin_sweep() elif path == "/admin/backup" and method == "POST": self._admin_backup(user_id) + elif path == "/admin/quotas" and method == "GET": + self._admin_quota_get(user_id, query) + elif path == "/admin/quotas" and method == "POST": + self._admin_quota_set(user_id) elif path == "/admin/push/dispatch" and method == "POST": self._push_dispatch_endpoint(user_id) elif path.startswith("/tasks/"): @@ -2017,6 +2038,8 @@ class CloudControlPlaneHandler(BaseHTTPRequestHandler): if not store.check_permission(user_id, project_id, "create_task"): self._send_json(403, {"error": "forbidden"}) return + if self._reject_if_over_quota(user_id): + return # 环境绑定：environment_id 非空时必须是本项目、已 ready # 的环境；creating/failed/archived 一律拒绝并把状态写进 # 报错（creating 时用户就知道该等 worker 准备好再建任务， @@ -2466,6 +2489,8 @@ class CloudControlPlaneHandler(BaseHTTPRequestHandler): user_id, task["project_id"], "create_task"): self._send_json(403, {"error": "forbidden"}) return + if self._reject_if_over_quota(user_id): + return message = body.get("message") if not message or not isinstance(message, str) \ or not message.strip(): @@ -2800,6 +2825,8 @@ class CloudControlPlaneHandler(BaseHTTPRequestHandler): user_id, task["project_id"], "approve_merge"): self._send_json(403, {"error": "forbidden"}) return + if self._reject_if_over_quota(user_id): + return if task["status"] not in ("failed", "interrupted", "merge-failed"): self._send_json(409, { "error": "only failed/interrupted tasks can be retried", @@ -2834,6 +2861,169 @@ class CloudControlPlaneHandler(BaseHTTPRequestHandler): store.audit(user_id, task["project_id"], "task_retried", task_id) self._send_json(200, task) + # -- 租户配额与计费 ---------------------------------------------- + # 计费口径：只算 cost_records 里 usd_micros 非空的实际成本 + # 行（estimated 行 usd 为 NULL，SUM 自动跳过，不算已花 + # 的钱）；归属按任务创建者 created_by，不按项目——一个 + # 用户在多个项目里花的钱都算进他自己的月度账单，项目维 + # 度的拆分由 _project_cost 另行提供。窗口是 UTC 自然月， + # 与 cost_records.day（UTC 日期）同口径。 + def _quota_month(self): + from datetime import datetime, timezone + return datetime.now(timezone.utc).strftime("%Y-%m") + + def _quota_status(self, user_id): + """返回 (quota_usd_micros|None, used_usd_micros, month)。 + quota 为 None 表示不限额：老用户与未设配额的用户默认 + 不限额，行为与过去一致，不会因为上了配额功能就突然 + 被挡住。""" + month = self._quota_month() + conn = self._db() + try: + _ensure_quotas_table(conn) + conn.commit() + # cost_records 由 CostStore 懒建：全新库上它还不存在， + # 直接 JOIN 会报 no such table——先实例化一次建表 + # （与 _project_cost 同款做法），再跑配额查询。 + _cost_store = CostStore(self.server.db_path) + try: + _cost_store._conn.close() + except Exception: + pass + row = conn.execute( + "SELECT quota_usd_micros FROM user_quotas" + " WHERE user_id = ?", + (user_id,)).fetchone() + quota = row[0] if row is not None else None + used = conn.execute( + "SELECT COALESCE(SUM(c.usd_micros), 0)" + " FROM cost_records c JOIN tasks t ON c.task_id = t.id" + " WHERE t.created_by = ? AND c.day LIKE ?", + (user_id, month + "%")).fetchone()[0] + return quota, int(used or 0), month + finally: + conn.close() + + def _quota_payload(self, user_id, quota, used, month): + return {"user_id": user_id, "month": month, + "quota_usd_micros": quota, "used_usd_micros": used, + "remaining_usd_micros": (quota - used) + if quota is not None else None} + + def _check_quota(self, user_id): + """配额准入：放行返回 None；超限返回 (quota, used, + month)，调用方一律回 402——这是钱用完了，不是请求 + 太快，所以不用 429。注意：只做建任务/分支/重跑时的 + 准入校验；跑起来的回合中途超限不会被掐断（worker + 侧没有按金额中途停转的钩子，那是下一步的事）。""" + quota, used, month = self._quota_status(user_id) + if quota is not None and used >= quota: + return quota, used, month + return None + + def _reject_if_over_quota(self, user_id): + blocked = self._check_quota(user_id) + if blocked is not None: + quota, used, month = blocked + self._send_json(402, {"error": "quota_exceeded", + "month": month, + "quota_usd_micros": quota, + "used_usd_micros": used}) + return True + return False + + def _is_global_maintainer(self, user_id): + """全局管理判定：沿用 _admin_backup 的口径——在任一 + 项目里是 maintainer/owner 即算（没有单设的全局管理 + 员表；没有任何项目角色的用户无权改任何人的配额）。""" + conn = self._db() + try: + row = conn.execute( + "SELECT COUNT(*) FROM project_members" + " WHERE user_id = ? AND role IN ('maintainer', 'owner')", + (user_id,)).fetchone() + return (row[0] if row is not None else 0) > 0 + finally: + conn.close() + + def _user_exists(self, target): + conn = self._db() + try: + return conn.execute( + "SELECT 1 FROM users WHERE id = ?", + (target,)).fetchone() is not None + finally: + conn.close() + + def _user_quota_self(self, user_id): + quota, used, month = self._quota_status(user_id) + self._send_json(200, + self._quota_payload(user_id, quota, used, month)) + + def _admin_quota_get(self, user_id, query): + if not self._is_global_maintainer(user_id): + self._send_json(403, {"error": "forbidden"}) + return + target = (query.get("user_id") or [None])[0] + if not target or not isinstance(target, str): + self._send_json(400, {"error": "user_id required"}) + return + if not self._user_exists(target): + self._send_json(404, {"error": "user not found"}) + return + quota, used, month = self._quota_status(target) + self._send_json(200, + self._quota_payload(target, quota, used, month)) + + def _admin_quota_set(self, user_id): + if not self._is_global_maintainer(user_id): + self._send_json(403, {"error": "forbidden"}) + return + body, err = self._read_json() + if err: + self._send_json(400, {"error": err}) + return + target = body.get("user_id") + if not target or not isinstance(target, str): + self._send_json(400, {"error": "user_id required"}) + return + if not self._user_exists(target): + self._send_json(404, {"error": "user not found"}) + return + quota = body.get("quota_usd_micros") + if quota is not None: + try: + quota = int(quota) + except (TypeError, ValueError): + self._send_json(400, { + "error": "quota_usd_micros must be an integer or null"}) + return + if quota < 0: + self._send_json(400, { + "error": "quota_usd_micros must be >= 0"}) + return + now = time.time() + conn = self._db() + try: + cur = conn.execute( + "UPDATE user_quotas SET quota_usd_micros = ?," + " updated_at = ? WHERE user_id = ?", + (quota, now, target)) + if cur.rowcount != 1: + conn.execute( + "INSERT INTO user_quotas(user_id, quota_usd_micros," + " updated_at) VALUES(?, ?, ?)", + (target, quota, now)) + conn.commit() + finally: + conn.close() + with _auth_store(self.server.db_path) as store: + store.audit(user_id, None, "quota_set", + "%s:%s" % (target, quota)) + quota, used, month = self._quota_status(target) + self._send_json(200, + self._quota_payload(target, quota, used, month)) + # -- worker 池 ----------------------------------------------------- def _workers_list(self, user_id): """列出注册 worker：任何已认证用户可读（worker 是基础设施 diff --git a/cloud/tests/test_task_quotas.py b/cloud/tests/test_task_quotas.py new file mode 100644 index 0000000..6539002 --- /dev/null +++ b/cloud/tests/test_task_quotas.py @@ -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()