"""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): 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): 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()