#!/usr/bin/env python3 # -*- coding: utf-8 -*- """new-api 额度池离线预置工具定向回归测试。""" from __future__ import annotations import importlib.util import sqlite3 import stat import tempfile import unittest from pathlib import Path from unittest import mock MODULE_PATH = Path(__file__).with_name("newapi_pool_provision.py") SPEC = importlib.util.spec_from_file_location("newapi_pool_provision", MODULE_PATH) assert SPEC and SPEC.loader MODULE = importlib.util.module_from_spec(SPEC) SPEC.loader.exec_module(MODULE) def token_state(*, token_key: str = "new-api-generated-secret-key", remain_quota: int = 6849315, status: int = 1, unlimited_quota: bool = False, expired_time: int = -1): """构造 token 权威状态，测试只覆盖与额度池有关的字段。""" return MODULE.TokenState( token_id=151, token_key=token_key, remain_quota=remain_quota, status=status, unlimited_quota=unlimited_quota, expired_time=expired_time, ) class ProvisionRequestTest(unittest.TestCase): """验证补建区间、宿主边界与敏感产物权限。""" def test_builds_exact_neice_051_to_054_range(self) -> None: names = MODULE.validate_provision_request(51, 4, "neice_", "mini-infra") self.assertEqual(names, ["neice_051", "neice_052", "neice_053", "neice_054"]) def test_accepts_canonical_runtime_hostname(self) -> None: names = MODULE.validate_provision_request(51, 4, "neice_", "minione-ubuntu-infra") self.assertEqual(names, ["neice_051", "neice_052", "neice_053", "neice_054"]) def test_rejects_non_mini_infra_host(self) -> None: with self.assertRaisesRegex(ValueError, "mini-infra"): MODULE.validate_provision_request(51, 4, "neice_", "mini-desktop") def test_rejects_empty_or_unsafe_ranges(self) -> None: for start, count, prefix in ((0, 4, "neice_"), (51, 0, "neice_"), (51, 4, "../neice_")): with self.subTest(start=start, count=count, prefix=prefix): with self.assertRaises(ValueError): MODULE.validate_provision_request(start, count, prefix, "mini-infra") def test_emit_files_are_owner_read_write_only(self) -> None: entry = { "newapi_user_id": 51, "newapi_token_id": 151, "newapi_token_key": "sk-secret-value-must-stay-in-transport", "grant_quota": 6849315, "quota_per_unit_snapshot": 500000, "usd_rate_snapshot": 7.3, "status": "FREE", "username": "neice_051", } with tempfile.TemporaryDirectory() as directory, mock.patch.object(MODULE, "datetime") as clock: clock.now.return_value.strftime.return_value = "20260722-010203" json_path, sql_path = MODULE.emit_outputs([entry], Path(directory), "neice_") self.assertEqual(stat.S_IMODE(json_path.stat().st_mode), 0o600) self.assertEqual(stat.S_IMODE(sql_path.stat().st_mode), 0o600) self.assertIn("-- pool_entry username=neice_051 newapi_user_id=51", sql_path.read_text()) def test_waits_for_api_write_to_become_visible(self) -> None: query = mock.Mock(side_effect=[None, None, 67]) with mock.patch.object(MODULE.time, "sleep") as sleep: value = MODULE.wait_for_db_visibility(query, attempts=4, interval_seconds=0.1) self.assertEqual(value, 67) self.assertEqual(query.call_count, 3) self.assertEqual(sleep.call_count, 2) def test_db_visibility_wait_is_bounded(self) -> None: query = mock.Mock(return_value=None) with mock.patch.object(MODULE.time, "sleep") as sleep: value = MODULE.wait_for_db_visibility(query, attempts=3, interval_seconds=0.1) self.assertIsNone(value) self.assertEqual(query.call_count, 3) self.assertEqual(sleep.call_count, 2) def test_postgres_attaches_sql_stdin_to_docker_exec(self) -> None: completed = mock.Mock(returncode=0, stdout="67\n", stderr="") with mock.patch.object(MODULE.subprocess, "run", return_value=completed) as run: uid = MODULE.Postgres("secret").query_user_id("neice_051") self.assertEqual(uid, 67) command = run.call_args.args[0] self.assertIn("-i", command) self.assertEqual(run.call_args.kwargs["input"], "SELECT id FROM users WHERE username='neice_051';") def test_repairs_quota_when_management_api_does_not_persist_it(self) -> None: pg = mock.Mock() pg.query_user_quota.side_effect = [0, 0, 0, 0, 0, 0, 6849315] with mock.patch.object(MODULE, "_http_json") as request, \ mock.patch.object(MODULE.time, "sleep"): result = MODULE.ensure_user_quota(pg, "root-token", 67, "neice_051", 6849315) self.assertEqual(result, "database-repair") request.assert_called_once() pg.set_user_quota.assert_called_once_with(67, 6849315) def test_reuses_verified_user_quota(self) -> None: pg = mock.Mock() pg.query_user_quota.return_value = 6849315 with mock.patch.object(MODULE, "_http_json") as request: result = MODULE.ensure_user_quota(pg, "root-token", 67, "neice_051", 6849315) self.assertEqual(result, "reused") request.assert_not_called() pg.set_user_quota.assert_not_called() class TokenStateTest(unittest.TestCase): """验证 token 权威余额与可用状态的修复、复核和 schema 兼容性。""" def assert_token_repaired(self, broken_state) -> None: pg = mock.Mock() pg.query_token.side_effect = [broken_state, token_state()] actual = MODULE.ensure_token_state( pg, 67, "neice_051", 6849315, created_in_this_run=True ) self.assertEqual(actual, token_state()) pg.repair_token.assert_called_once_with( 151, 6849315, created_in_this_run=True ) def test_repairs_newly_created_disabled_token(self) -> None: self.assert_token_repaired(token_state(status=0)) def test_repairs_newly_created_zero_token_balance(self) -> None: self.assert_token_repaired(token_state(remain_quota=0)) def test_repairs_newly_created_unlimited_token(self) -> None: self.assert_token_repaired(token_state(unlimited_quota=True)) def test_repairs_newly_created_expired_token(self) -> None: self.assert_token_repaired(token_state(expired_time=1_800_000_000)) def test_default_rerun_never_refills_consumed_existing_token(self) -> None: pg = mock.Mock() pg.query_token.return_value = token_state(remain_quota=0) with self.assertRaisesRegex(RuntimeError, "既有 token.*禁止修复或充值"): MODULE.ensure_token_state(pg, 67, "neice_051", 6849315) pg.repair_token.assert_not_called() def test_default_rerun_never_reactivates_existing_disabled_token(self) -> None: pg = mock.Mock() pg.query_token.return_value = token_state(status=0) with self.assertRaisesRegex(RuntimeError, "既有 token.*禁止修复或充值"): MODULE.ensure_token_state(pg, 67, "neice_051", 6849315) pg.repair_token.assert_not_called() def test_direct_repair_is_denied_without_new_token_proof(self) -> None: pg = MODULE.Postgres.__new__(MODULE.Postgres) pg._run = mock.Mock() with self.assertRaisesRegex(RuntimeError, "只允许修复本次新建 token"): pg.repair_token(151, 6849315) pg._run.assert_not_called() def test_rejects_repair_when_authoritative_token_state_stays_invalid(self) -> None: pg = mock.Mock() pg.query_token.side_effect = [token_state(status=0), token_state(status=0)] with self.assertRaisesRegex(RuntimeError, "token 权威状态修复后仍不一致"): MODULE.ensure_token_state( pg, 67, "neice_051", 6849315, created_in_this_run=True ) def test_token_sql_is_compatible_with_sqlite_schema(self) -> None: connection = sqlite3.connect(":memory:") connection.execute( "CREATE TABLE tokens (" "id INTEGER PRIMARY KEY, user_id INTEGER, name TEXT, key TEXT, " "remain_quota INTEGER, status INTEGER, unlimited_quota BOOLEAN, expired_time INTEGER)" ) connection.execute( "INSERT INTO tokens VALUES (151, 67, 'neice_051', 'secret-key', 0, 0, 1, 123)" ) pg = MODULE.Postgres.__new__(MODULE.Postgres) def run_sql(sql: str, operation: str = "query") -> str: cursor = connection.execute(sql) if cursor.description is None: connection.commit() return "" return "\n".join("|".join(str(value) for value in row) for row in cursor.fetchall()) pg._run = run_sql before = pg.query_token(67, "neice_051") self.assertEqual(before, token_state( token_key="secret-key", remain_quota=0, status=0, unlimited_quota=True, expired_time=123 )) pg.repair_token(151, 6849315, created_in_this_run=True) self.assertEqual(pg.query_token(67, "neice_051"), token_state(token_key="secret-key")) connection.close() class ProvisionSecurityTest(unittest.TestCase): """验证随机管理凭据、同名 token 幂等与日志脱敏。""" def build_existing_user(self, access_token: str = "") -> mock.Mock: pg = mock.Mock() pg.query_user_id.return_value = 67 pg.query_access_token.return_value = access_token pg.query_user_quota.return_value = 6849315 return pg def test_new_access_tokens_are_random_for_same_username(self) -> None: generated = iter(("random-access-token-first-123456", "random-access-token-second-654321")) written_tokens = [] with mock.patch.object(MODULE.secrets, "token_urlsafe", side_effect=lambda size: next(generated)), \ mock.patch.object(MODULE, "_http_json"): for token_key in ("gateway-key-first", "gateway-key-second"): pg = self.build_existing_user() pg.query_token.side_effect = [None, MODULE.TokenState( 151, token_key, 6849315, 1, False, -1 ), MODULE.TokenState(151, token_key, 6849315, 1, False, -1), MODULE.TokenState(151, token_key, 6849315, 1, False, -1), MODULE.TokenState(151, token_key, 6849315, 1, False, -1)] MODULE.provision_one(pg, "root-token", "neice_051", 6849315, 500000, 7.3, False) written_tokens.append(pg.set_access_token.call_args.args[1]) self.assertNotEqual(written_tokens[0], written_tokens[1]) self.assertNotIn("neice_051", written_tokens) def test_rerun_reuses_existing_user_and_token(self) -> None: pg = self.build_existing_user("existing-random-access-token-123") pg.query_token.return_value = token_state() with mock.patch.object(MODULE, "_http_json") as request, \ mock.patch.object(MODULE.secrets, "token_urlsafe") as random_token: result = MODULE.provision_one( pg, "root-token", "neice_051", 6849315, 500000, 7.3, False ) self.assertEqual(result["newapi_token_id"], 151) self.assertEqual(result["newapi_token_key"], "new-api-generated-secret-key") request.assert_not_called() random_token.assert_not_called() pg.set_access_token.assert_not_called() pg.repair_token.assert_not_called() def test_rotates_legacy_predictable_access_token(self) -> None: legacy_token = MODULE.legacy_access_token("neice_051") pg = self.build_existing_user(legacy_token) pg.query_token.return_value = token_state() with mock.patch.object(MODULE.secrets, "token_urlsafe", return_value="rotated-random-token-123456789"): MODULE.provision_one(pg, "root-token", "neice_051", 6849315, 500000, 7.3, False) pg.set_access_token.assert_called_once_with(67, "rotated-random-token-123456789") def test_logs_do_not_reveal_access_or_gateway_token_fragments(self) -> None: access_token = "random-access-token-sensitive-123456" gateway_key = "gateway-token-sensitive-654321" pg = self.build_existing_user() state = MODULE.TokenState(151, gateway_key, 6849315, 1, False, -1) pg.query_token.side_effect = [None, state, state, state] with mock.patch.object(MODULE.secrets, "token_urlsafe", return_value=access_token), \ mock.patch.object(MODULE, "_http_json"), \ self.assertLogs("newapi_pool", level="INFO") as captured: MODULE.provision_one(pg, "root-token", "neice_051", 6849315, 500000, 7.3, False) logs = "\n".join(captured.output) for secret in (access_token, gateway_key): self.assertNotIn(secret, logs) self.assertNotIn(secret[:4], logs) self.assertNotIn(secret[-4:], logs) def test_existing_consumed_token_fails_before_user_quota_write(self) -> None: pg = self.build_existing_user(MODULE.legacy_access_token("neice_051")) pg.query_token.return_value = token_state(remain_quota=0) with mock.patch.object(MODULE.secrets, "token_urlsafe") as random_token: with self.assertRaisesRegex(RuntimeError, "禁止修复或充值"): MODULE.provision_one( pg, "root-token", "neice_051", 6849315, 500000, 7.3, False ) pg.repair_token.assert_not_called() pg.set_access_token.assert_not_called() pg.set_user_quota.assert_not_called() random_token.assert_not_called() if __name__ == "__main__": unittest.main()