mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-01 17:00:21 +08:00
perf: 优化请求鉴权链路并批量化统计/调度查询
- 为 Pipeline/Context 增加按需读取请求体能力，支持 async 懒加载 JSON body - 为 chat/cli/video/claude/openai-cli 适配器关闭默认预读，减少无效 body 读取与超时风险 - 将本地登录、JWT 用户加载、API Key 鉴权迁移到线程池隔离会话执行，避免阻塞事件循环 - 为 API Key 鉴权返回结构化余额结果，并在主请求会话中重新绑定 user/api_key 后再校验状态、过期和锁定信息 - 为 management/user token 前缀认证引入独立会话与结果回绑，避免跨会话对象写入失效 - 为 Usage 余额检查补充结构化返回，统一透出 remaining 与欠费/不可用文案映射 - 为用户与管理端活跃请求查询增加 maintain_status 开关，避免轮询指定 id 时误触发状态修复 - 重写 user_me usage 汇总逻辑，支持 group_by=None 的粗粒度聚合 - 修正 provider 维度成功率与平均响应时间统计，基于 success_count 和成功响应耗时汇总计算 - 前端 Usage 轮询由 setInterval 改为串行 setTimeout，避免并发轮询叠加 - 为 StatsAggregator 增加按本地日期批量计算百分位能力，替代逐天 fan-out 查询 - 为混合统计查询合并连续实时日期区间，并批量读取 StatsDaily，减少逐日查询次数 - 为用户日统计增加批量聚合入口，替代逐用户循环聚合 - 为系统配置导出改用 selectinload 预加载 provider 关联数据，减少 N+1 查询 - 为管理员用户列表增加钱包批量查询，避免逐用户回表 - 为调度器增加 provider 轻量引用预过滤，先按 allowed_providers 缩小范围再加载完整 provider 图 - 为 CandidateBuilder 增加 provider refs/provider_ids 查询能力，保留分页顺序 - 为模型缓存增加 provider_model_mappings 索引缓存与 model_mappings 规则缓存，减少重复全量扫描 - 为请求候选中间态改为 flush/batch commit，降低 pending/streaming 状态切换的事务往返 - 为钱包访问结果补充 balance_snapshot，并抽取余额快照复用逻辑 - 补充 pipeline、auth、admin users、user_me usage、stats aggregator、model cache、 scheduler、wallet、request candidate 等回归与契约测试
This commit is contained in:
AAEE86
93
tests/api/test_admin_user_routes.py
Normal file
93
tests/api/test_admin_user_routes.py
Normal file
@@ -0,0 +1,93 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from datetime import datetime, timezone
|
||||
from types import SimpleNamespace
|
||||
from typing import Any
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
import pytest
|
||||
from fastapi import FastAPI
|
||||
from fastapi.testclient import TestClient
|
||||
|
||||
from src.api.admin.users.routes import router as admin_users_router
|
||||
from src.database import get_db
|
||||
|
||||
|
||||
def _build_admin_users_app(db: MagicMock, monkeypatch: pytest.MonkeyPatch) -> TestClient:
|
||||
app = FastAPI()
|
||||
app.include_router(admin_users_router)
|
||||
app.dependency_overrides[get_db] = lambda: db
|
||||
|
||||
async def _fake_pipeline_run(
|
||||
*, adapter: Any, http_request: object, db: MagicMock, mode: object
|
||||
) -> Any:
|
||||
_ = http_request, mode
|
||||
context = SimpleNamespace(
|
||||
db=db,
|
||||
request=SimpleNamespace(state=SimpleNamespace()),
|
||||
user=SimpleNamespace(id="admin-1"),
|
||||
ensure_json_body=lambda: {},
|
||||
add_audit_metadata=lambda **_: None,
|
||||
)
|
||||
return await adapter.handle(context)
|
||||
|
||||
monkeypatch.setattr("src.api.admin.users.routes.pipeline.run", _fake_pipeline_run)
|
||||
return TestClient(app)
|
||||
|
||||
|
||||
def test_list_users_uses_wallet_batch_lookup(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
db = MagicMock()
|
||||
client = _build_admin_users_app(db, monkeypatch)
|
||||
now = datetime.now(timezone.utc)
|
||||
users = [
|
||||
SimpleNamespace(
|
||||
id="user-1",
|
||||
email="u1@example.com",
|
||||
username="user1",
|
||||
role=SimpleNamespace(value="user"),
|
||||
allowed_providers=None,
|
||||
allowed_api_formats=None,
|
||||
allowed_models=None,
|
||||
is_active=True,
|
||||
created_at=now,
|
||||
updated_at=now,
|
||||
last_login_at=None,
|
||||
),
|
||||
SimpleNamespace(
|
||||
id="user-2",
|
||||
email="u2@example.com",
|
||||
username="user2",
|
||||
role=SimpleNamespace(value="admin"),
|
||||
allowed_providers=None,
|
||||
allowed_api_formats=None,
|
||||
allowed_models=None,
|
||||
is_active=True,
|
||||
created_at=now,
|
||||
updated_at=None,
|
||||
last_login_at=None,
|
||||
),
|
||||
]
|
||||
wallets_by_user_id = {
|
||||
"user-1": SimpleNamespace(limit_mode="unlimited"),
|
||||
}
|
||||
batch_getter = MagicMock(return_value=wallets_by_user_id)
|
||||
|
||||
monkeypatch.setattr(
|
||||
"src.api.admin.users.routes.UserService.list_users", lambda *_a, **_k: users
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
"src.api.admin.users.routes.WalletService.get_wallets_by_user_ids",
|
||||
batch_getter,
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
"src.api.admin.users.routes.WalletService.get_wallet",
|
||||
lambda *_a, **_k: (_ for _ in ()).throw(AssertionError("不应回退到逐个钱包查询")),
|
||||
)
|
||||
|
||||
response = client.get("/api/admin/users")
|
||||
|
||||
assert response.status_code == 200
|
||||
assert response.json()[0]["unlimited"] is True
|
||||
assert response.json()[1]["unlimited"] is False
|
||||
batch_getter.assert_called_once()
|
||||
assert batch_getter.call_args.args[1] == ["user-1", "user-2"]
|
||||
@@ -8,54 +8,139 @@ API Pipeline 测试
|
||||
"""
|
||||
|
||||
from datetime import datetime, timezone
|
||||
from typing import Any
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import pytest
|
||||
from fastapi import HTTPException
|
||||
|
||||
from src.api.base.adapter import ApiMode
|
||||
from src.api.base.pipeline import ApiRequestPipeline
|
||||
from src.core.enums import UserRole
|
||||
from src.core.modules.hooks import AUTH_TOKEN_PREFIX_AUTHENTICATORS
|
||||
|
||||
|
||||
class TestPipelineBalanceCalculation:
|
||||
"""测试 Pipeline 余额计算"""
|
||||
"""Balance calculation tests for Pipeline."""
|
||||
|
||||
@pytest.fixture
|
||||
def pipeline(self) -> ApiRequestPipeline:
|
||||
return ApiRequestPipeline()
|
||||
|
||||
def test_calculate_balance_remaining_with_balance(self, pipeline: ApiRequestPipeline) -> None:
|
||||
"""测试有限制钱包时计算剩余余额"""
|
||||
@pytest.mark.asyncio
|
||||
async def test_calculate_balance_remaining_with_balance(
|
||||
self, pipeline: ApiRequestPipeline
|
||||
) -> None:
|
||||
"""Returns remaining balance for limited wallets."""
|
||||
mock_user = MagicMock()
|
||||
mock_db = MagicMock()
|
||||
mock_user.id = "user-123"
|
||||
|
||||
with patch(
|
||||
"src.api.base.pipeline.WalletService.get_balance_snapshot",
|
||||
return_value=70.0,
|
||||
):
|
||||
remaining = pipeline._calculate_balance_remaining(mock_db, mock_user)
|
||||
thread_db = MagicMock()
|
||||
db_user = MagicMock()
|
||||
thread_db.query.return_value.filter.return_value.first.return_value = db_user
|
||||
|
||||
with patch("src.api.base.pipeline.create_session", return_value=thread_db):
|
||||
with patch(
|
||||
"src.api.base.pipeline.WalletService.get_balance_snapshot",
|
||||
return_value=70.0,
|
||||
):
|
||||
remaining = await pipeline._calculate_balance_remaining_async(mock_user)
|
||||
|
||||
assert remaining == 70.0
|
||||
thread_db.close.assert_called_once()
|
||||
|
||||
def test_calculate_balance_remaining_unlimited(self, pipeline: ApiRequestPipeline) -> None:
|
||||
"""测试无限制钱包时返回 None"""
|
||||
@pytest.mark.asyncio
|
||||
async def test_calculate_balance_remaining_unlimited(
|
||||
self, pipeline: ApiRequestPipeline
|
||||
) -> None:
|
||||
"""Returns None for unlimited wallets."""
|
||||
mock_user = MagicMock()
|
||||
mock_db = MagicMock()
|
||||
mock_user.id = "user-123"
|
||||
|
||||
with patch(
|
||||
"src.api.base.pipeline.WalletService.get_balance_snapshot",
|
||||
return_value=None,
|
||||
thread_db = MagicMock()
|
||||
db_user = MagicMock()
|
||||
thread_db.query.return_value.filter.return_value.first.return_value = db_user
|
||||
|
||||
with patch("src.api.base.pipeline.create_session", return_value=thread_db):
|
||||
with patch(
|
||||
"src.api.base.pipeline.WalletService.get_balance_snapshot",
|
||||
return_value=None,
|
||||
):
|
||||
remaining = await pipeline._calculate_balance_remaining_async(mock_user)
|
||||
|
||||
assert remaining is None
|
||||
thread_db.close.assert_called_once()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_calculate_balance_remaining_none_user(
|
||||
self, pipeline: ApiRequestPipeline
|
||||
) -> None:
|
||||
"""Returns None when user is missing."""
|
||||
remaining = await pipeline._calculate_balance_remaining_async(None)
|
||||
|
||||
assert remaining is None
|
||||
|
||||
|
||||
class TestPipelineRunModes:
|
||||
"""Returns remaining balance for limited wallets."""
|
||||
|
||||
@pytest.fixture
|
||||
def pipeline(self) -> ApiRequestPipeline:
|
||||
return ApiRequestPipeline()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_run_management_mode_skips_balance_calculation(
|
||||
self, pipeline: ApiRequestPipeline
|
||||
) -> None:
|
||||
"""Management mode skips balance calculation."""
|
||||
mock_request = MagicMock()
|
||||
mock_request.method = "GET"
|
||||
mock_request.url.path = "/api/admin/tokens"
|
||||
mock_request.state = MagicMock()
|
||||
|
||||
mock_db = MagicMock()
|
||||
mock_user = MagicMock()
|
||||
mock_user.id = "admin-123"
|
||||
mock_token = MagicMock()
|
||||
mock_token.id = "mt-123"
|
||||
|
||||
mock_adapter = MagicMock()
|
||||
mock_adapter.name = "test-adapter"
|
||||
mock_adapter.authorize = MagicMock(return_value=None)
|
||||
mock_response = MagicMock()
|
||||
mock_response.status_code = 200
|
||||
mock_adapter.handle = AsyncMock(return_value=mock_response)
|
||||
|
||||
mock_context = MagicMock()
|
||||
mock_context.db = mock_db
|
||||
mock_context.request = mock_request
|
||||
|
||||
with patch.object(
|
||||
pipeline,
|
||||
"_authenticate_management",
|
||||
new_callable=AsyncMock,
|
||||
return_value=(mock_user, mock_token),
|
||||
):
|
||||
remaining = pipeline._calculate_balance_remaining(mock_db, mock_user)
|
||||
with patch(
|
||||
"src.api.base.pipeline.ApiRequestContext.build",
|
||||
return_value=mock_context,
|
||||
):
|
||||
with patch.object(
|
||||
pipeline,
|
||||
"_calculate_balance_remaining_async",
|
||||
new_callable=AsyncMock,
|
||||
) as mock_balance:
|
||||
with patch.object(pipeline, "_record_audit_event"):
|
||||
response = await pipeline.run(
|
||||
mock_adapter,
|
||||
mock_request,
|
||||
mock_db,
|
||||
mode=ApiMode.MANAGEMENT,
|
||||
)
|
||||
|
||||
assert remaining is None
|
||||
|
||||
def test_calculate_balance_remaining_none_user(self, pipeline: ApiRequestPipeline) -> None:
|
||||
"""测试用户为 None 时返回 None"""
|
||||
mock_db = MagicMock()
|
||||
remaining = pipeline._calculate_balance_remaining(mock_db, None)
|
||||
|
||||
assert remaining is None
|
||||
assert response == mock_response
|
||||
assert mock_context.management_token == mock_token
|
||||
mock_balance.assert_not_called()
|
||||
|
||||
|
||||
class TestPipelineAuditLogging:
|
||||
@@ -213,7 +298,8 @@ class TestPipelineAuthentication:
|
||||
def pipeline(self) -> ApiRequestPipeline:
|
||||
return ApiRequestPipeline()
|
||||
|
||||
def test_authenticate_client_missing_key(self, pipeline: ApiRequestPipeline) -> None:
|
||||
@pytest.mark.asyncio
|
||||
async def test_authenticate_client_missing_key(self, pipeline: ApiRequestPipeline) -> None:
|
||||
"""测试缺少 API Key 时抛出异常"""
|
||||
mock_request = MagicMock()
|
||||
mock_request.headers = {}
|
||||
@@ -226,12 +312,13 @@ class TestPipelineAuthentication:
|
||||
mock_adapter.extract_api_key = MagicMock(return_value=None)
|
||||
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
pipeline._authenticate_client(mock_request, mock_db, mock_adapter)
|
||||
await pipeline._authenticate_client(mock_request, mock_db, mock_adapter)
|
||||
|
||||
assert exc_info.value.status_code == 401
|
||||
assert "API密钥" in exc_info.value.detail
|
||||
|
||||
def test_authenticate_client_invalid_key(self, pipeline: ApiRequestPipeline) -> None:
|
||||
@pytest.mark.asyncio
|
||||
async def test_authenticate_client_invalid_key(self, pipeline: ApiRequestPipeline) -> None:
|
||||
"""测试无效的 API Key"""
|
||||
mock_request = MagicMock()
|
||||
mock_request.headers = {"Authorization": "Bearer sk-invalid"}
|
||||
@@ -245,15 +332,17 @@ class TestPipelineAuthentication:
|
||||
|
||||
with patch.object(
|
||||
pipeline.auth_service,
|
||||
"authenticate_api_key",
|
||||
"authenticate_api_key_threadsafe",
|
||||
new_callable=AsyncMock,
|
||||
return_value=None,
|
||||
):
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
pipeline._authenticate_client(mock_request, mock_db, mock_adapter)
|
||||
await pipeline._authenticate_client(mock_request, mock_db, mock_adapter)
|
||||
|
||||
assert exc_info.value.status_code == 401
|
||||
|
||||
def test_authenticate_client_balance_exceeded(self, pipeline: ApiRequestPipeline) -> None:
|
||||
@pytest.mark.asyncio
|
||||
async def test_authenticate_client_balance_exceeded(self, pipeline: ApiRequestPipeline) -> None:
|
||||
"""测试余额不足时抛出异常"""
|
||||
mock_user = MagicMock()
|
||||
mock_user.id = "user-123"
|
||||
@@ -268,28 +357,198 @@ class TestPipelineAuthentication:
|
||||
mock_request.state = MagicMock()
|
||||
|
||||
mock_db = MagicMock()
|
||||
db_user = MagicMock()
|
||||
db_user.id = "user-123"
|
||||
db_user.is_active = True
|
||||
db_user.is_deleted = False
|
||||
db_api_key = MagicMock()
|
||||
db_api_key.id = "key-123"
|
||||
db_api_key.user_id = "user-123"
|
||||
db_api_key.is_active = True
|
||||
db_api_key.is_locked = False
|
||||
db_api_key.is_standalone = False
|
||||
db_api_key.expires_at = None
|
||||
user_query = MagicMock()
|
||||
user_query.filter.return_value.first.return_value = db_user
|
||||
api_key_query = MagicMock()
|
||||
api_key_query.filter.return_value.first.return_value = db_api_key
|
||||
mock_db.query.side_effect = [user_query, api_key_query]
|
||||
|
||||
mock_adapter = MagicMock()
|
||||
mock_adapter.extract_api_key = MagicMock(return_value="sk-test")
|
||||
|
||||
with patch.object(
|
||||
pipeline.auth_service,
|
||||
"authenticate_api_key",
|
||||
return_value=(mock_user, mock_api_key),
|
||||
"authenticate_api_key_threadsafe",
|
||||
new_callable=AsyncMock,
|
||||
return_value=MagicMock(
|
||||
user=mock_user,
|
||||
api_key=mock_api_key,
|
||||
access_allowed=False,
|
||||
balance_remaining=0.0,
|
||||
),
|
||||
):
|
||||
with patch.object(
|
||||
pipeline.usage_service,
|
||||
"check_request_balance",
|
||||
return_value=(False, "余额不足"),
|
||||
):
|
||||
with patch(
|
||||
"src.api.base.pipeline.WalletService.get_balance_snapshot",
|
||||
return_value=0.0,
|
||||
):
|
||||
from src.core.exceptions import BalanceInsufficientException
|
||||
from src.core.exceptions import BalanceInsufficientException
|
||||
|
||||
with pytest.raises(BalanceInsufficientException):
|
||||
pipeline._authenticate_client(mock_request, mock_db, mock_adapter)
|
||||
with pytest.raises(BalanceInsufficientException):
|
||||
await pipeline._authenticate_client(mock_request, mock_db, mock_adapter)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_authenticate_client_requery_detects_inactive_user(
|
||||
self, pipeline: ApiRequestPipeline
|
||||
) -> None:
|
||||
mock_user = MagicMock()
|
||||
mock_user.id = "user-123"
|
||||
mock_api_key = MagicMock()
|
||||
mock_api_key.id = "key-123"
|
||||
|
||||
mock_request = MagicMock()
|
||||
mock_request.headers = {"Authorization": "Bearer sk-test"}
|
||||
mock_request.url.path = "/v1/messages"
|
||||
mock_request.state = MagicMock()
|
||||
|
||||
mock_adapter = MagicMock()
|
||||
mock_adapter.extract_api_key = MagicMock(return_value="sk-test")
|
||||
|
||||
db_user = MagicMock()
|
||||
db_user.id = "user-123"
|
||||
db_user.is_active = False
|
||||
db_user.is_deleted = False
|
||||
db_api_key = MagicMock()
|
||||
db_api_key.id = "key-123"
|
||||
db_api_key.user_id = "user-123"
|
||||
db_api_key.is_active = True
|
||||
db_api_key.is_locked = False
|
||||
db_api_key.is_standalone = False
|
||||
db_api_key.expires_at = None
|
||||
|
||||
mock_db = MagicMock()
|
||||
user_query = MagicMock()
|
||||
user_query.filter.return_value.first.return_value = db_user
|
||||
api_key_query = MagicMock()
|
||||
api_key_query.filter.return_value.first.return_value = db_api_key
|
||||
mock_db.query.side_effect = [user_query, api_key_query]
|
||||
|
||||
with patch.object(
|
||||
pipeline.auth_service,
|
||||
"authenticate_api_key_threadsafe",
|
||||
new_callable=AsyncMock,
|
||||
return_value=MagicMock(
|
||||
user=mock_user,
|
||||
api_key=mock_api_key,
|
||||
access_allowed=True,
|
||||
balance_remaining=10.0,
|
||||
),
|
||||
):
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
await pipeline._authenticate_client(mock_request, mock_db, mock_adapter)
|
||||
|
||||
assert exc_info.value.status_code == 401
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_authenticate_client_requery_detects_locked_key(
|
||||
self, pipeline: ApiRequestPipeline
|
||||
) -> None:
|
||||
mock_user = MagicMock()
|
||||
mock_user.id = "user-123"
|
||||
mock_api_key = MagicMock()
|
||||
mock_api_key.id = "key-123"
|
||||
|
||||
mock_request = MagicMock()
|
||||
mock_request.headers = {"Authorization": "Bearer sk-test"}
|
||||
mock_request.url.path = "/v1/messages"
|
||||
mock_request.state = MagicMock()
|
||||
|
||||
mock_adapter = MagicMock()
|
||||
mock_adapter.extract_api_key = MagicMock(return_value="sk-test")
|
||||
|
||||
db_user = MagicMock()
|
||||
db_user.id = "user-123"
|
||||
db_user.is_active = True
|
||||
db_user.is_deleted = False
|
||||
db_api_key = MagicMock()
|
||||
db_api_key.id = "key-123"
|
||||
db_api_key.user_id = "user-123"
|
||||
db_api_key.is_active = True
|
||||
db_api_key.is_locked = True
|
||||
db_api_key.is_standalone = False
|
||||
db_api_key.expires_at = None
|
||||
|
||||
mock_db = MagicMock()
|
||||
user_query = MagicMock()
|
||||
user_query.filter.return_value.first.return_value = db_user
|
||||
api_key_query = MagicMock()
|
||||
api_key_query.filter.return_value.first.return_value = db_api_key
|
||||
mock_db.query.side_effect = [user_query, api_key_query]
|
||||
|
||||
with patch.object(
|
||||
pipeline.auth_service,
|
||||
"authenticate_api_key_threadsafe",
|
||||
new_callable=AsyncMock,
|
||||
return_value=MagicMock(
|
||||
user=mock_user,
|
||||
api_key=mock_api_key,
|
||||
access_allowed=True,
|
||||
balance_remaining=10.0,
|
||||
),
|
||||
):
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
await pipeline._authenticate_client(mock_request, mock_db, mock_adapter)
|
||||
|
||||
assert exc_info.value.status_code == 403
|
||||
assert "锁定" in str(exc_info.value.detail)
|
||||
|
||||
|
||||
class TestPipelineTokenPrefixAuth:
|
||||
"""Tests token-prefix auth isolation."""
|
||||
|
||||
@pytest.fixture
|
||||
def pipeline(self) -> ApiRequestPipeline:
|
||||
return ApiRequestPipeline()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_try_token_prefix_auth_uses_isolated_session(
|
||||
self, pipeline: ApiRequestPipeline
|
||||
) -> None:
|
||||
mock_request = MagicMock()
|
||||
mock_request.headers = {}
|
||||
mock_request.client = MagicMock(host="127.0.0.1")
|
||||
|
||||
route_db = MagicMock()
|
||||
auth_db = MagicMock()
|
||||
mock_user = MagicMock()
|
||||
mock_token = MagicMock()
|
||||
|
||||
async def authenticate(db: Any, token: str, client_ip: str) -> tuple[Any, Any]:
|
||||
assert db is auth_db
|
||||
assert token == "ae_test"
|
||||
assert client_ip == "127.0.0.1"
|
||||
return mock_user, mock_token
|
||||
|
||||
with patch("src.api.base.pipeline.create_session", return_value=auth_db):
|
||||
with patch("src.utils.request_utils.get_client_ip", return_value="127.0.0.1"):
|
||||
with patch("src.core.modules.hooks.get_hook_dispatcher") as mock_get_dispatcher:
|
||||
dispatcher = MagicMock()
|
||||
dispatcher.dispatch = AsyncMock(
|
||||
return_value=[
|
||||
{
|
||||
"prefix": "ae_",
|
||||
"module": "management_tokens",
|
||||
"authenticate": authenticate,
|
||||
}
|
||||
]
|
||||
)
|
||||
mock_get_dispatcher.return_value = dispatcher
|
||||
|
||||
result = await pipeline._try_token_prefix_auth(
|
||||
"ae_test", mock_request, route_db
|
||||
)
|
||||
|
||||
assert result == (mock_user, mock_token)
|
||||
dispatcher.dispatch.assert_awaited_once_with(AUTH_TOKEN_PREFIX_AUTHENTICATORS)
|
||||
auth_db.expunge.assert_any_call(mock_user)
|
||||
auth_db.expunge.assert_any_call(mock_token)
|
||||
auth_db.close.assert_called_once()
|
||||
|
||||
|
||||
class TestPipelineAdminAuth:
|
||||
|
||||
175
tests/api/test_user_me_usage_routes.py
Normal file
175
tests/api/test_user_me_usage_routes.py
Normal file
@@ -0,0 +1,175 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
import pytest
|
||||
|
||||
from src.api.user_me.routes import GetUsageAdapter
|
||||
from src.core.enums import UserRole
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_usage_adapter_uses_coarse_summary_grouping(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
db = MagicMock()
|
||||
query = MagicMock()
|
||||
count_query = MagicMock()
|
||||
count_query.scalar.return_value = 0
|
||||
query.outerjoin.return_value = query
|
||||
query.filter.return_value = query
|
||||
query.with_entities.return_value = count_query
|
||||
query.options.return_value = query
|
||||
query.order_by.return_value = query
|
||||
query.offset.return_value = query
|
||||
query.limit.return_value = query
|
||||
query.all.return_value = []
|
||||
db.query.return_value = query
|
||||
|
||||
summary_getter = MagicMock(
|
||||
return_value=[
|
||||
{
|
||||
"provider": "provider-a",
|
||||
"model": "gpt-4o",
|
||||
"requests": 2,
|
||||
"input_tokens": 10,
|
||||
"output_tokens": 5,
|
||||
"total_tokens": 15,
|
||||
"total_cost_usd": 1.5,
|
||||
"actual_total_cost_usd": 1.2,
|
||||
"success_count": 2,
|
||||
"success_response_time_sum_ms": 1000.0,
|
||||
"success_response_time_count": 2,
|
||||
},
|
||||
{
|
||||
"provider": "pending",
|
||||
"model": "gpt-4o",
|
||||
"requests": 99,
|
||||
"input_tokens": 999,
|
||||
"output_tokens": 999,
|
||||
"total_tokens": 1998,
|
||||
"total_cost_usd": 9.9,
|
||||
"actual_total_cost_usd": 9.9,
|
||||
"success_count": 0,
|
||||
"success_response_time_sum_ms": 0.0,
|
||||
"success_response_time_count": 0,
|
||||
},
|
||||
]
|
||||
)
|
||||
monkeypatch.setattr("src.api.user_me.routes.UsageService.get_usage_summary", summary_getter)
|
||||
monkeypatch.setattr("src.api.user_me.routes.WalletService.get_wallet", lambda *_a, **_k: None)
|
||||
monkeypatch.setattr(
|
||||
"src.api.user_me.routes.WalletService.serialize_wallet_summary",
|
||||
lambda _wallet: {"limit_mode": "finite"},
|
||||
)
|
||||
|
||||
adapter = GetUsageAdapter(time_range=None, limit=20, offset=0)
|
||||
context = SimpleNamespace(
|
||||
db=db,
|
||||
user=SimpleNamespace(id="user-1", role=UserRole.USER),
|
||||
request=SimpleNamespace(state=SimpleNamespace()),
|
||||
)
|
||||
|
||||
result = await adapter.handle(context)
|
||||
|
||||
assert result["total_requests"] == 2
|
||||
assert result["total_tokens"] == 15
|
||||
assert result["summary_by_model"] == [
|
||||
{
|
||||
"model": "gpt-4o",
|
||||
"requests": 2,
|
||||
"input_tokens": 10,
|
||||
"output_tokens": 5,
|
||||
"total_tokens": 15,
|
||||
"total_cost_usd": 1.5,
|
||||
}
|
||||
]
|
||||
assert "total_actual_cost" not in result
|
||||
assert result["summary_by_provider"] == [
|
||||
{
|
||||
"provider": "provider-a",
|
||||
"requests": 2,
|
||||
"total_tokens": 15,
|
||||
"total_cost_usd": 1.5,
|
||||
"success_rate": 100.0,
|
||||
"avg_response_time_ms": 500.0,
|
||||
}
|
||||
]
|
||||
assert summary_getter.call_args.kwargs["group_by"] is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_usage_adapter_provider_success_rate_uses_success_count(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
db = MagicMock()
|
||||
query = MagicMock()
|
||||
count_query = MagicMock()
|
||||
count_query.scalar.return_value = 0
|
||||
query.outerjoin.return_value = query
|
||||
query.filter.return_value = query
|
||||
query.with_entities.return_value = count_query
|
||||
query.options.return_value = query
|
||||
query.order_by.return_value = query
|
||||
query.offset.return_value = query
|
||||
query.limit.return_value = query
|
||||
query.all.return_value = []
|
||||
db.query.return_value = query
|
||||
|
||||
summary_getter = MagicMock(
|
||||
return_value=[
|
||||
{
|
||||
"provider": "provider-a",
|
||||
"model": "gpt-4o",
|
||||
"requests": 3,
|
||||
"input_tokens": 30,
|
||||
"output_tokens": 15,
|
||||
"total_tokens": 45,
|
||||
"total_cost_usd": 4.5,
|
||||
"actual_total_cost_usd": 4.5,
|
||||
"success_count": 2,
|
||||
"success_response_time_sum_ms": 600.0,
|
||||
"success_response_time_count": 2,
|
||||
},
|
||||
{
|
||||
"provider": "provider-a",
|
||||
"model": "gpt-4.1",
|
||||
"requests": 1,
|
||||
"input_tokens": 10,
|
||||
"output_tokens": 5,
|
||||
"total_tokens": 15,
|
||||
"total_cost_usd": 1.5,
|
||||
"actual_total_cost_usd": 1.5,
|
||||
"success_count": 0,
|
||||
"success_response_time_sum_ms": 0.0,
|
||||
"success_response_time_count": 0,
|
||||
},
|
||||
]
|
||||
)
|
||||
monkeypatch.setattr("src.api.user_me.routes.UsageService.get_usage_summary", summary_getter)
|
||||
monkeypatch.setattr("src.api.user_me.routes.WalletService.get_wallet", lambda *_a, **_k: None)
|
||||
monkeypatch.setattr(
|
||||
"src.api.user_me.routes.WalletService.serialize_wallet_summary",
|
||||
lambda _wallet: {"limit_mode": "finite"},
|
||||
)
|
||||
|
||||
adapter = GetUsageAdapter(time_range=None, limit=20, offset=0)
|
||||
context = SimpleNamespace(
|
||||
db=db,
|
||||
user=SimpleNamespace(id="user-1", role=UserRole.USER),
|
||||
request=SimpleNamespace(state=SimpleNamespace()),
|
||||
)
|
||||
|
||||
result = await adapter.handle(context)
|
||||
|
||||
assert result["summary_by_provider"] == [
|
||||
{
|
||||
"provider": "provider-a",
|
||||
"requests": 4,
|
||||
"total_tokens": 60,
|
||||
"total_cost_usd": 6.0,
|
||||
"success_rate": 50.0,
|
||||
"avg_response_time_ms": 300.0,
|
||||
}
|
||||
]
|
||||
Reference in New Issue
Block a user