mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-02 09:20:22 +08:00
76a6d0ce8e
- 新增 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 等测试
142 lines
4.0 KiB
Python
142 lines
4.0 KiB
Python
"""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
|