310 lines
14 KiB
Python
Executable File
310 lines
14 KiB
Python
Executable File
#!/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()
|