mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-02 17:30:23 +08:00
feat(pool): 增加 OAuth 账号池管理功能
- 新增 pool manager / strategy / health_policy / cost_tracker / redis_ops 等核心模块 - 新增 pool admin API 路由与 schemas - 新增 OAuth 账号类型解析 (oauth_plan) - 前端增加 PoolManagement 页面、PoolConfigDialog、PoolImportDialog、PoolStatusCard 组件 - 补充 pool config / cost tracker / health policy / manager / strategy / trace 等测试
This commit is contained in:
141
tests/services/test_pool_cost_tracker.py
Normal file
141
tests/services/test_pool_cost_tracker.py
Normal file
@@ -0,0 +1,141 @@
|
||||
"""Tests for pool_cost_tracker.py — rolling window cost tracking."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from unittest.mock import AsyncMock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
from src.services.provider.pool.config import PoolConfig
|
||||
from src.services.provider.pool.cost_tracker import (
|
||||
get_window_usage,
|
||||
is_approaching_limit,
|
||||
is_at_limit,
|
||||
record_usage,
|
||||
)
|
||||
|
||||
PID = "provider-test"
|
||||
KID = "key-test"
|
||||
|
||||
|
||||
@pytest.fixture()
|
||||
def config() -> PoolConfig:
|
||||
return PoolConfig(
|
||||
cost_limit_per_key_tokens=10000,
|
||||
cost_soft_threshold_percent=80,
|
||||
cost_window_seconds=18000,
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture()
|
||||
def config_no_limit() -> PoolConfig:
|
||||
return PoolConfig(cost_limit_per_key_tokens=None)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# record_usage
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_record_usage_calls_redis(config: PoolConfig) -> None:
|
||||
with patch(
|
||||
"src.services.provider.pool.redis_ops.add_cost_entry",
|
||||
new_callable=AsyncMock,
|
||||
) as mock_add:
|
||||
await record_usage(PID, KID, 500, config)
|
||||
|
||||
mock_add.assert_called_once_with(PID, KID, 500, 18000)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_record_usage_skips_zero_tokens(config: PoolConfig) -> None:
|
||||
with patch(
|
||||
"src.services.provider.pool.redis_ops.add_cost_entry",
|
||||
new_callable=AsyncMock,
|
||||
) as mock_add:
|
||||
await record_usage(PID, KID, 0, config)
|
||||
|
||||
mock_add.assert_not_called()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_record_usage_skips_when_no_limit(config_no_limit: PoolConfig) -> None:
|
||||
with patch(
|
||||
"src.services.provider.pool.redis_ops.add_cost_entry",
|
||||
new_callable=AsyncMock,
|
||||
) as mock_add:
|
||||
await record_usage(PID, KID, 500, config_no_limit)
|
||||
|
||||
mock_add.assert_not_called()
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# is_at_limit
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_is_at_limit_true(config: PoolConfig) -> None:
|
||||
with patch(
|
||||
"src.services.provider.pool.redis_ops.get_cost_window_total",
|
||||
new_callable=AsyncMock,
|
||||
return_value=10000,
|
||||
):
|
||||
assert await is_at_limit(PID, KID, config) is True
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_is_at_limit_false(config: PoolConfig) -> None:
|
||||
with patch(
|
||||
"src.services.provider.pool.redis_ops.get_cost_window_total",
|
||||
new_callable=AsyncMock,
|
||||
return_value=5000,
|
||||
):
|
||||
assert await is_at_limit(PID, KID, config) is False
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_is_at_limit_always_false_when_no_limit(config_no_limit: PoolConfig) -> None:
|
||||
assert await is_at_limit(PID, KID, config_no_limit) is False
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# is_approaching_limit
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_is_approaching_limit_true(config: PoolConfig) -> None:
|
||||
# 80% of 10000 = 8000
|
||||
with patch(
|
||||
"src.services.provider.pool.redis_ops.get_cost_window_total",
|
||||
new_callable=AsyncMock,
|
||||
return_value=8500,
|
||||
):
|
||||
assert await is_approaching_limit(PID, KID, config) is True
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_is_approaching_limit_false(config: PoolConfig) -> None:
|
||||
with patch(
|
||||
"src.services.provider.pool.redis_ops.get_cost_window_total",
|
||||
new_callable=AsyncMock,
|
||||
return_value=5000,
|
||||
):
|
||||
assert await is_approaching_limit(PID, KID, config) is False
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# get_window_usage
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_window_usage(config: PoolConfig) -> None:
|
||||
with patch(
|
||||
"src.services.provider.pool.redis_ops.get_cost_window_total",
|
||||
new_callable=AsyncMock,
|
||||
return_value=4200,
|
||||
):
|
||||
assert await get_window_usage(PID, KID, config) == 4200
|
||||
Reference in New Issue
Block a user