Files
host 8c27407b4a fix: 登录限流下开通突发不再硬失败(节流 + 429 重试)
网关登录端点 20 次/分钟/源 IP,100 路突发开通首测仅 19
成功:同步重试波(2/4/8/16/32 秒)把线程推到下一窗口
开头重测也只到 39 成功。本批两层:
- 进程内登录节流:60 秒窗口最多 15 次,满了等最早名额
  滑出窗口再进(留 5 次给其它进程与人工登录);注册是
  best-effort,失败仍有 worker 首任务懒开通兜底,故不定
  时提前放弃;
- 429 退避重试保留:用户登录与管理员登录遇 429 时
  2/4/8/16/32 秒 + 抖动重试,最多 6 次,兜底跨进程竞争;
  管理员登录最终失败才清共享缓存。
测试:新增 100 并发单测(预算临时抬高,只验并发正确性)
与 3 个 429 重试单测、1 个节流窗口单测(fake clock);
单测 setUp/tearDown 清空进程级节流名额,免跨测累积拖慢
套件;全量 cloud 800(4 跳过)、host 69、deploy 19 全绿;
真网关 100 路突发重测 100/100 成功、墙钟 362 秒、临时
用户 100/100 已清。残留:硬保证需共享 Redis 节流,控制
面目前没有 Redis 客户端,未做。
2026-10-27 23:09:22 +08:00

532 lines
21 KiB
Python
Raw Permalink Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
"""sub2api 网关接入单测:管理 API 客户端 + 幂等开通 + provider 覆盖。
不跑真实模型、不联网:HTTP 经 request 可调用对象注入假实现;
开通走内存 SQLite;host/sandbox 侧只断言命令与 env 构造。
"""
import os
import sqlite3
import sys
import threading
import unittest
from pathlib import Path
from unittest import mock
try:
from cloud import sub2api
from cloud.sub2api import Sub2apiAdmin, ensure_user_gateway_key
except ImportError: # discover -s cloud/tests 时把项目根目录补进 sys.path
sys.path.insert(0, str(Path(__file__).resolve().parents[2]))
from cloud import sub2api
from cloud.sub2api import Sub2apiAdmin, ensure_user_gateway_key
def _conn():
conn = sqlite3.connect(":memory:")
conn.row_factory = sqlite3.Row
return conn
class FakeApi:
"""按 (method, path) 录调用、按表回包的假 sub2api。"""
def __init__(self):
self.calls = []
self._lock = threading.Lock()
self.admin_token = "adm-tok"
self.user_token = "usr-tok"
self.groups = [{"id": 2, "name": "codex-system"}]
self.next_user_id = 10
self.next_key_id = 20
self.fail_refresh_once = False
self.balance_calls = []
self.fail_balance = False
self.fail_create_user = False
self.create_user_hook = None
self.login_429_times = 0
self.admin_login_429_times = 0
def __call__(self, method, url, body=None, token=None):
with self._lock:
return self._handle(method, url, body, token)
def _handle(self, method, url, body=None, token=None):
path = url.split("28084", 1)[-1] if "28084" in url else url
self.calls.append((method, path, body, token))
if path == "/api/v1/auth/login" and body.get("email") == "adm@x":
if self.admin_login_429_times > 0:
self.admin_login_429_times -= 1
return 429, {"message":
"Too many requests, please try again later"}
return 200, {"data": {"access_token": self.admin_token,
"refresh_token": "adm-ref",
"expires_in": 3600}}
if path == "/api/v1/auth/login":
if self.login_429_times > 0:
self.login_429_times -= 1
return 429, {"message":
"Too many requests, please try again later"}
return 200, {"data": {"access_token": self.user_token,
"refresh_token": "usr-ref",
"expires_in": 3600}}
if path == "/api/v1/auth/refresh":
return 200, {"data": {"access_token": "new-tok",
"refresh_token": "new-ref",
"expires_in": 3600}}
if path.startswith("/api/v1/admin/users/") \
and path.endswith("/balance") and method == "POST":
self.balance_calls.append(body)
if self.fail_balance:
return 400, {"message": "bad balance"}
return 200, {"data": {"balance": body.get("balance")}}
if path == "/api/v1/admin/users" and method == "POST":
self.next_user_id += 1
if self.create_user_hook is not None:
self.create_user_hook(self.next_user_id)
if self.fail_create_user:
return 400, {"message": "email already exists"}
return 200, {"data": {"id": self.next_user_id}}
if path == "/api/v1/admin/groups" and method == "GET":
return 200, {"data": self.groups}
if path == "/api/v1/keys" and method == "POST":
self.next_key_id += 1
return 200, {"data": {"id": self.next_key_id,
"key": "sk-probe-%d" % self.next_key_id}}
if path == "/v1/models" and method == "GET":
return 200, {"data": [{"id": "gpt-5"}, {"id": "claude-x"}]}
return 404, {"message": "not found"}
class TestSub2apiAdmin(unittest.TestCase):
def setUp(self):
self.api = FakeApi()
self.admin = Sub2apiAdmin("http://x:28084", "adm@x", "pw",
request=self.api)
def test_create_user_and_group_lookup(self):
uid = self.admin.create_user("u@x", "pw", "u1")
self.assertEqual(uid, 11)
self.assertEqual(self.admin.resolve_group_id(), 2)
def test_401_refreshes_and_retries(self):
# 第一次管理调用先 401:注入一个先 401 后放行的 request。
calls = {"n": 0}
real = self.api
def flaky(method, url, body=None, token=None):
calls["n"] += 1
if calls["n"] == 2:
return 401, {"message": "expired"}
return real(method, url, body, token)
admin = Sub2apiAdmin("http://x:28084", "adm@x", "pw",
request=flaky)
uid = admin.create_user("u@x", "pw", "u1")
self.assertEqual(uid, 11)
paths = [c[1] for c in self.api.calls]
self.assertIn("/api/v1/auth/refresh", paths)
class TestEnsureGatewayKey(unittest.TestCase):
def setUp(self):
self.api = FakeApi()
self.admin = Sub2apiAdmin("http://x:28084", "adm@x", "pw",
request=self.api)
# 节流名额是进程级共享状态:每个单测从空开始,避免跨测累积。
sub2api._login_slots.clear()
self.env = {"SUB2API_BASE_URL": "http://x:28084",
"SUB2API_GROUP_NAME": "codex-system"}
self.old = {k: os.environ.get(k) for k in self.env}
os.environ.update(self.env)
def tearDown(self):
sub2api._login_slots.clear()
for k, v in self.old.items():
if v is None:
os.environ.pop(k, None)
else:
os.environ[k] = v
def test_provision_then_idempotent(self):
conn = _conn()
key = ensure_user_gateway_key(conn, "user-abc", admin=self.admin,
request=self.api)
self.assertTrue(key.startswith("sk-probe-"))
calls_before = len(self.api.calls)
key2 = ensure_user_gateway_key(conn, "user-abc", admin=self.admin,
request=self.api)
self.assertEqual(key, key2)
# 第二次是纯读库,不再调管理 API。
self.assertEqual(len(self.api.calls), calls_before)
row = conn.execute(
"SELECT sub2api_user_id, refresh_token FROM user_gateway_keys"
" WHERE user_id = ?", ("user-abc",)).fetchone()
self.assertEqual(row[0], 11)
self.assertEqual(row[1], "usr-ref")
def test_list_models(self):
models = sub2api.list_gateway_models("http://x:28084", "sk-1",
request=self.api)
self.assertEqual(models, ["gpt-5", "claude-x"])
def test_provision_grants_opening_balance(self):
conn = _conn()
key = ensure_user_gateway_key(conn, "user-bal", admin=self.admin,
request=self.api)
self.assertTrue(key.startswith("sk-probe-"))
self.assertEqual(len(self.api.balance_calls), 1)
self.assertEqual(self.api.balance_calls[0]["balance"], 10.0)
self.assertEqual(self.api.balance_calls[0]["operation"], "set")
def test_provision_balance_failure_cleans_up(self):
self.api.fail_balance = True
conn = _conn()
with self.assertRaises(sub2api.Sub2apiError):
ensure_user_gateway_key(conn, "user-fail", admin=self.admin,
request=self.api)
paths = [(m, p) for (m, p, _b, _t) in self.api.calls]
self.assertIn(("DELETE", "/api/v1/keys/21"), paths)
self.assertIn(("DELETE", "/api/v1/admin/users/11"), paths)
row = conn.execute(
"SELECT COUNT(*) FROM user_gateway_keys"
" WHERE user_id = ?", ("user-fail",)).fetchone()
self.assertEqual(row[0], 0)
def test_provision_zero_opening_balance_skips_grant(self):
old = os.environ.get("CLOUD_SUB2API_OPENING_BALANCE_USD")
os.environ["CLOUD_SUB2API_OPENING_BALANCE_USD"] = "0"
try:
conn = _conn()
ensure_user_gateway_key(conn, "user-zero", admin=self.admin,
request=self.api)
finally:
if old is None:
os.environ.pop("CLOUD_SUB2API_OPENING_BALANCE_USD", None)
else:
os.environ["CLOUD_SUB2API_OPENING_BALANCE_USD"] = old
self.assertEqual(self.api.balance_calls, [])
def test_provision_race_reuses_winner_row(self):
conn = _conn()
def hook(_new_id):
conn.execute(
"INSERT INTO user_gateway_keys(user_id, sub2api_user_id,"
" api_key, refresh_token, created_at, updated_at)"
" VALUES(?, ?, ?, ?, ?, ?)",
("user-race", 99, "sk-race", "ref-race", 0, 0))
conn.commit()
self.api.create_user_hook = hook
self.api.fail_create_user = True
key = ensure_user_gateway_key(conn, "user-race", admin=self.admin,
request=self.api)
self.assertEqual(key, "sk-race")
def test_provision_race_no_row_reraises(self):
self.api.fail_create_user = True
conn = _conn()
with mock.patch("time.sleep", return_value=None):
with self.assertRaises(sub2api.Sub2apiError):
ensure_user_gateway_key(conn, "user-race2",
admin=self.admin, request=self.api)
row = conn.execute(
"SELECT COUNT(*) FROM user_gateway_keys"
" WHERE user_id = ?", ("user-race2",)).fetchone()
self.assertEqual(row[0], 0)
def _admin_env(self):
old = {k: os.environ.get(k) for k in
("SUB2API_ADMIN_EMAIL", "SUB2API_ADMIN_PASSWORD")}
os.environ["SUB2API_ADMIN_EMAIL"] = "adm@x"
os.environ["SUB2API_ADMIN_PASSWORD"] = "pw"
sub2api._SHARED_ADMIN.clear()
return old
def _restore_admin_env(self, old):
for k, v in old.items():
if v is None:
os.environ.pop(k, None)
else:
os.environ[k] = v
sub2api._SHARED_ADMIN.clear()
def _admin_login_count(self):
return len([
c for c in self.api.calls
if c[0] == "POST" and c[1] == "/api/v1/auth/login"
and (c[2] or {}).get("email") == "adm@x"])
def test_shared_admin_session_reuses_login(self):
old = self._admin_env()
try:
conn1 = _conn()
key1 = ensure_user_gateway_key(conn1, "user-s1",
request=self.api)
self.assertTrue(key1.startswith("sk-probe-"))
conn2 = _conn()
key2 = ensure_user_gateway_key(conn2, "user-s2",
request=self.api)
self.assertTrue(key2.startswith("sk-probe-"))
# 同一进程内第二次开通不再管理员登录(只剩首次那一次)。
self.assertEqual(self._admin_login_count(), 1)
finally:
self._restore_admin_env(old)
def test_shared_admin_session_expired_refetches(self):
old = self._admin_env()
sub2api._SHARED_ADMIN[("http://x:28084", "adm@x")] = (
"old-tok", "old-ref", 0)
try:
conn = _conn()
key = ensure_user_gateway_key(conn, "user-s3",
request=self.api)
self.assertTrue(key.startswith("sk-probe-"))
self.assertEqual(self._admin_login_count(), 1)
finally:
self._restore_admin_env(old)
def test_shared_admin_session_concurrent_burst_logs_in_once(self):
old = self._admin_env()
barrier = threading.Barrier(10)
results = []
errors = []
def run(i):
try:
barrier.wait(timeout=10)
conn = _conn()
key = ensure_user_gateway_key(
conn, "user-c%d" % i, request=self.api)
results.append(key)
except Exception as exc: # noqa: BLE001 - 汇总后断言
errors.append(exc)
try:
threads = [threading.Thread(target=run, args=(i,))
for i in range(10)]
for t in threads:
t.start()
for t in threads:
t.join(timeout=30)
self.assertEqual(errors, [])
self.assertEqual(len(results), 10)
self.assertTrue(all(k.startswith("sk-probe-")
for k in results))
# 突发冷缓存:10 路并发也只登录一次。
self.assertEqual(self._admin_login_count(), 1)
finally:
self._restore_admin_env(old)
def test_shared_admin_session_concurrent_100(self):
# 本测只验并发正确性(各 key 唯一、管理员登录 1 次),不验节流:
# 把预算临时抬高,免得 101 次登录被 15 次/分钟的节流拖到 7 分钟。
orig_budget = sub2api._LOGIN_BUDGET_PER_MINUTE
sub2api._LOGIN_BUDGET_PER_MINUTE = 1000
old = self._admin_env()
barrier = threading.Barrier(100)
results = []
errors = []
def run(i):
try:
barrier.wait(timeout=30)
conn = _conn()
key = ensure_user_gateway_key(
conn, "user-h%d" % i, request=self.api)
results.append(key)
except Exception as exc: # noqa: BLE001 - 汇总后断言
errors.append(exc)
try:
threads = [threading.Thread(target=run, args=(i,))
for i in range(100)]
for t in threads:
t.start()
for t in threads:
t.join(timeout=120)
self.assertEqual(errors, [])
self.assertEqual(len(results), 100)
self.assertTrue(all(k.startswith("sk-probe-")
for k in results))
self.assertEqual(len(set(results)), 100)
# 100 路突发冷缓存也只登录一次。
self.assertEqual(self._admin_login_count(), 1)
finally:
sub2api._LOGIN_BUDGET_PER_MINUTE = orig_budget
self._restore_admin_env(old)
def test_user_login_retries_on_429_then_succeeds(self):
self.api.login_429_times = 2
with mock.patch("time.sleep", return_value=None):
token, refresh = sub2api.user_login(
"http://x:28084", "u@x", "pw", request=self.api)
self.assertEqual(token, "usr-tok")
self.assertEqual(refresh, "usr-ref")
logins = [c for c in self.api.calls
if c[1] == "/api/v1/auth/login"
and (c[2] or {}).get("email") != "adm@x"]
self.assertEqual(len(logins), 3)
def test_user_login_429_exhausted_raises(self):
self.api.login_429_times = 100
with mock.patch("time.sleep", return_value=None):
with self.assertRaises(sub2api.Sub2apiError):
sub2api.user_login("http://x:28084", "u@x", "pw",
request=self.api)
logins = [c for c in self.api.calls
if c[1] == "/api/v1/auth/login"
and (c[2] or {}).get("email") != "adm@x"]
self.assertEqual(len(logins), 6)
def test_admin_login_retries_on_429(self):
self.api.admin_login_429_times = 2
admin = Sub2apiAdmin("http://x:28084", "adm@x", "pw",
request=self.api)
with mock.patch("time.sleep", return_value=None):
uid = admin.create_user("u@x", "pw", "u1")
self.assertEqual(uid, 11)
self.assertEqual(self._admin_login_count(), 3)
def test_login_slot_throttle_waits_for_window(self):
orig_budget = sub2api._LOGIN_BUDGET_PER_MINUTE
orig_window = sub2api._LOGIN_WINDOW_SECONDS
sub2api._LOGIN_BUDGET_PER_MINUTE = 2
sub2api._LOGIN_WINDOW_SECONDS = 60.0
sub2api._login_slots.clear()
now = [1000.0]
sleeps = []
def fake_time():
return now[0]
def fake_sleep(seconds):
sleeps.append(seconds)
now[0] += seconds
try:
with mock.patch("time.time", side_effect=fake_time), \
mock.patch("time.sleep", side_effect=fake_sleep):
sub2api._acquire_login_slot()
sub2api._acquire_login_slot()
sub2api._acquire_login_slot()
self.assertEqual(sleeps, [60.0])
finally:
sub2api._LOGIN_BUDGET_PER_MINUTE = orig_budget
sub2api._LOGIN_WINDOW_SECONDS = orig_window
sub2api._login_slots.clear()
class TestProviderCommand(unittest.TestCase):
def test_server_command_sub2api(self):
sys.path.insert(0, str(Path(__file__).resolve().parents[2]))
import host
old = os.environ.get("SUB2API_BASE_URL")
os.environ["SUB2API_BASE_URL"] = "http://172.17.0.1:28084"
try:
cmd = host._server_command("/bin/codex", provider="sub2api")
finally:
if old is None:
os.environ.pop("SUB2API_BASE_URL", None)
else:
os.environ["SUB2API_BASE_URL"] = old
self.assertIn("app-server", cmd[-1])
joined = " ".join(cmd)
self.assertIn('model_provider="sub2api"', joined)
self.assertIn("http://172.17.0.1:28084/v1", joined)
self.assertIn('env_key="SUB2API_API_KEY"', joined)
def test_server_command_sub2api_needs_base_url(self):
sys.path.insert(0, str(Path(__file__).resolve().parents[2]))
import host
old = os.environ.pop("SUB2API_BASE_URL", None)
try:
with self.assertRaises(RuntimeError):
host._server_command("/bin/codex", provider="sub2api")
finally:
if old is not None:
os.environ["SUB2API_BASE_URL"] = old
def test_server_command_unknown_provider(self):
sys.path.insert(0, str(Path(__file__).resolve().parents[2]))
import host
with self.assertRaises(ValueError):
host._server_command("/bin/codex", provider="nope")
class TestSandboxProviderFlag(unittest.TestCase):
def test_base_command_flag(self):
from cloud.sandbox import LocalExecutor, DockerExecutor
local = LocalExecutor("/tmp/host.py")
cmd = local._base_command(model_provider="sub2api")
self.assertIn("--model-provider", cmd)
self.assertIn("sub2api", cmd)
docker = DockerExecutor("/tmp/host.py")
inner = docker._inner_host_command("dispatch",
model_provider="sub2api")
self.assertIn("--model-provider", inner)
def test_run_local_env_merge_and_loud_fail(self):
from cloud.sandbox import LocalExecutor
seen = {}
def runner(cmd, env=None):
seen["env"] = env
return object()
local = LocalExecutor("/tmp/host.py", runner=runner)
local._run_local(["x"], {"SUB2API_API_KEY": "sk-1"})
self.assertEqual(seen["env"]["SUB2API_API_KEY"], "sk-1")
def cmd_only(cmd):
return object()
local2 = LocalExecutor("/tmp/host.py", runner=cmd_only)
with self.assertRaises(RuntimeError):
local2._run_local(["x"], {"SUB2API_API_KEY": "sk-1"})
class TestDockerSecretEnv(unittest.TestCase):
def _docker(self, runner):
from cloud.sandbox import DockerExecutor
ex = DockerExecutor("/tmp/host.py", runner=runner, env_key=None)
ex._check_docker = lambda: None
return ex
def test_secret_name_only_in_argv(self):
from cloud.sandbox import DockerExecutor
seen = {}
def runner(cmd, timeout=None, env=None):
seen["cmd"] = cmd
seen["env"] = env
return object()
ex = self._docker(runner)
task = {"id": "abc", "prompt": "hi"}
cmd = ex.build_run_command(
task, ["python3", "/app/host.py", "status"],
extra_env={"SUB2API_API_KEY": "sk-secret",
"HTTP_PROXY": "http://x:8888"})
self.assertIn("SUB2API_API_KEY", cmd)
self.assertNotIn("SUB2API_API_KEY=sk-secret", cmd)
self.assertIn("HTTP_PROXY=http://x:8888", cmd)
env = ex._docker_cli_env({"SUB2API_API_KEY": "sk-secret"})
self.assertEqual(env["SUB2API_API_KEY"], "sk-secret")
self.assertIsNone(ex._docker_cli_env({"HTTP_PROXY": "http://x:1"}))
self.assertIsNone(ex._docker_cli_env(None))
def test_call_runner_env_loud_fail(self):
from cloud.sandbox import DockerExecutor
def cmd_only(cmd, timeout=None):
return object()
ex = self._docker(cmd_only)
with self.assertRaises(RuntimeError):
ex._call_runner(["docker"], env={"SUB2API_API_KEY": "sk-1"})
if __name__ == "__main__":
unittest.main()