codex-app-server/cloud/tests/test_sub2api.py

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