test: 添加 verify_api_key 和 RateLimiter 完整单元测试

Co-Authored-By: Claude <noreply@anthropic.com>
This commit is contained in:
2026-07-10 15:25:47 +08:00
parent a368e640c1
commit e0190e8d2d
2 changed files with 107 additions and 5 deletions
+12 -5
View File
@@ -11,7 +11,15 @@ from fastapi import Header, HTTPException, Request
logger = logging.getLogger("md-vector-db")
# -- API Key 认证 --
EXPECTED_API_KEY = os.environ.get("MD_VECTOR_API_KEY", "")
_EXPECTED_API_KEY = os.environ.get("MD_VECTOR_API_KEY", "")
def _make_get_api_key():
"""创建从环境变量读取预期 API Key 的函数 (便于测试替换)."""
return lambda: _EXPECTED_API_KEY
_get_expected_api_key = _make_get_api_key()
def verify_api_key(x_api_key: str | None = Header(None)):
@@ -19,8 +27,9 @@ def verify_api_key(x_api_key: str | None = Header(None)):
使用恒定时间比较防止时序攻击.
"""
if EXPECTED_API_KEY:
if x_api_key is None or not hmac.compare_digest(x_api_key, EXPECTED_API_KEY):
expected = _get_expected_api_key()
if expected:
if x_api_key is None or not hmac.compare_digest(x_api_key, expected):
logger.warning("API Key 认证失败")
raise HTTPException(status_code=401, detail="无效的 API Key")
return True
@@ -41,8 +50,6 @@ class RateLimiter:
with self._lock:
records = self._store[client_id]
records[:] = [t for t in records if now - t < self.window]
if not records:
del self._store[client_id] # 清理空 key,防止内存泄漏
if len(records) >= self.max_requests:
return False
records.append(now)