host
706d5a8ca4
首版共享会话在缓存未命中时各线程各自登录：突发冷缓存下 N 路并发开通会在网关刷出 N 次 bcrypt。改为锁内双重检查 登录（RLock 可重入，_login 成功后写回同一缓存）：第一路 登录，其余线程等锁后直接播种。实例本身仍不跨线程，401 时照常走刷新->重登录自愈。 测试：FakeApi 加锁（假实现自身线程安全，失败指向真代码）； 新增 10 线程同起并发单测，断言 10 路全成功且管理员登录 恰好 1 次；全量 cloud 795（4 跳过）、host 69、deploy 19 全绿；真网关 10 路突发实测 0 错误、墙钟 0.774s、临时号 已清。
421 lines
16 KiB
Python
421 lines
16 KiB
Python
"""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
|
||
|
||
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":
|
||
return 200, {"data": {"access_token": self.admin_token,
|
||
"refresh_token": "adm-ref",
|
||
"expires_in": 3600}}
|
||
if path == "/api/v1/auth/login":
|
||
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)
|
||
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):
|
||
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)
|
||
|
||
|
||
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()
|