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:
fawney19
274
tests/services/test_pool_manager.py
Normal file
274
tests/services/test_pool_manager.py
Normal file
@@ -0,0 +1,274 @@
|
||||
"""Tests for pool_manager.py — candidate reordering, success/error hooks."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import AsyncMock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
from src.services.provider.pool.config import PoolConfig
|
||||
from src.services.provider.pool.manager import PoolManager
|
||||
|
||||
|
||||
def _make_candidate(key_id: str, *, is_skipped: bool = False) -> SimpleNamespace:
|
||||
return SimpleNamespace(
|
||||
key=SimpleNamespace(id=key_id),
|
||||
is_skipped=is_skipped,
|
||||
skip_reason=None,
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture()
|
||||
def pool() -> PoolManager:
|
||||
return PoolManager("provider-1", PoolConfig())
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# reorder_candidates
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_reorder_empty_candidates(pool: PoolManager) -> None:
|
||||
result = await pool.reorder_candidates("sess-1", [])
|
||||
assert result == []
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_reorder_sticky_hit_moves_to_front(pool: PoolManager) -> None:
|
||||
c1 = _make_candidate("key-1")
|
||||
c2 = _make_candidate("key-2")
|
||||
c3 = _make_candidate("key-3")
|
||||
|
||||
with (
|
||||
patch(
|
||||
"src.services.provider.pool.redis_ops.get_sticky_binding",
|
||||
new_callable=AsyncMock,
|
||||
return_value="key-2",
|
||||
),
|
||||
patch(
|
||||
"src.services.provider.pool.redis_ops.batch_get_cooldowns",
|
||||
new_callable=AsyncMock,
|
||||
return_value={"key-1": None, "key-2": None, "key-3": None},
|
||||
),
|
||||
patch(
|
||||
"src.services.provider.pool.redis_ops.get_lru_scores",
|
||||
new_callable=AsyncMock,
|
||||
return_value={"key-1": 100.0, "key-2": 200.0, "key-3": 50.0},
|
||||
),
|
||||
):
|
||||
result = await pool.reorder_candidates("sess-1", [c1, c2, c3])
|
||||
|
||||
assert result[0].key.id == "key-2" # sticky hit first
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_reorder_cooldown_keys_are_skipped(pool: PoolManager) -> None:
|
||||
c1 = _make_candidate("key-1")
|
||||
c2 = _make_candidate("key-2")
|
||||
|
||||
with (
|
||||
patch(
|
||||
"src.services.provider.pool.redis_ops.get_sticky_binding",
|
||||
new_callable=AsyncMock,
|
||||
return_value=None,
|
||||
),
|
||||
patch(
|
||||
"src.services.provider.pool.redis_ops.batch_get_cooldowns",
|
||||
new_callable=AsyncMock,
|
||||
return_value={"key-1": "rate_limited_429", "key-2": None},
|
||||
),
|
||||
patch(
|
||||
"src.services.provider.pool.redis_ops.get_lru_scores",
|
||||
new_callable=AsyncMock,
|
||||
return_value={},
|
||||
),
|
||||
):
|
||||
result = await pool.reorder_candidates(None, [c1, c2])
|
||||
|
||||
# key-1 should be skipped
|
||||
assert c1.is_skipped is True
|
||||
assert "cooldown" in (c1.skip_reason or "")
|
||||
# key-2 first in available
|
||||
assert result[0].key.id == "key-2"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_reorder_cost_exhausted_keys_are_skipped() -> None:
|
||||
pool = PoolManager(
|
||||
"provider-1",
|
||||
PoolConfig(cost_limit_per_key_tokens=1000, lru_enabled=False),
|
||||
)
|
||||
c1 = _make_candidate("key-1")
|
||||
c2 = _make_candidate("key-2")
|
||||
|
||||
with (
|
||||
patch(
|
||||
"src.services.provider.pool.redis_ops.get_sticky_binding",
|
||||
new_callable=AsyncMock,
|
||||
return_value=None,
|
||||
),
|
||||
patch(
|
||||
"src.services.provider.pool.redis_ops.batch_get_cooldowns",
|
||||
new_callable=AsyncMock,
|
||||
return_value={"key-1": None, "key-2": None},
|
||||
),
|
||||
patch(
|
||||
"src.services.provider.pool.redis_ops.batch_get_cost_totals",
|
||||
new_callable=AsyncMock,
|
||||
return_value={"key-1": 1500, "key-2": 200},
|
||||
),
|
||||
):
|
||||
result = await pool.reorder_candidates(None, [c1, c2])
|
||||
|
||||
assert c1.is_skipped is True
|
||||
assert "cost" in (c1.skip_reason or "")
|
||||
assert result[0].key.id == "key-2"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_reorder_lru_sorts_least_recently_used_first(
|
||||
pool: PoolManager,
|
||||
) -> None:
|
||||
c1 = _make_candidate("key-1")
|
||||
c2 = _make_candidate("key-2")
|
||||
c3 = _make_candidate("key-3")
|
||||
|
||||
with (
|
||||
patch(
|
||||
"src.services.provider.pool.redis_ops.get_sticky_binding",
|
||||
new_callable=AsyncMock,
|
||||
return_value=None,
|
||||
),
|
||||
patch(
|
||||
"src.services.provider.pool.redis_ops.batch_get_cooldowns",
|
||||
new_callable=AsyncMock,
|
||||
return_value={"key-1": None, "key-2": None, "key-3": None},
|
||||
),
|
||||
patch(
|
||||
"src.services.provider.pool.redis_ops.get_lru_scores",
|
||||
new_callable=AsyncMock,
|
||||
return_value={"key-1": 300.0, "key-2": 100.0, "key-3": 200.0},
|
||||
),
|
||||
):
|
||||
result = await pool.reorder_candidates(None, [c1, c2, c3])
|
||||
|
||||
# Least recently used (lowest score) first
|
||||
available = [r for r in result if not r.is_skipped]
|
||||
assert available[0].key.id == "key-2"
|
||||
assert available[1].key.id == "key-3"
|
||||
assert available[2].key.id == "key-1"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# on_request_success
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_on_request_success_binds_sticky_and_touches_lru(
|
||||
pool: PoolManager,
|
||||
) -> None:
|
||||
with (
|
||||
patch(
|
||||
"src.services.provider.pool.redis_ops.set_sticky_binding",
|
||||
new_callable=AsyncMock,
|
||||
) as mock_sticky,
|
||||
patch(
|
||||
"src.services.provider.pool.redis_ops.touch_lru",
|
||||
new_callable=AsyncMock,
|
||||
) as mock_lru,
|
||||
patch(
|
||||
"src.services.provider.pool.redis_ops.add_cost_entry",
|
||||
new_callable=AsyncMock,
|
||||
) as mock_cost,
|
||||
):
|
||||
await pool.on_request_success(session_uuid="sess-1", key_id="key-1", tokens_used=0)
|
||||
|
||||
mock_sticky.assert_called_once_with("provider-1", "sess-1", "key-1", 3600)
|
||||
mock_lru.assert_called_once_with("provider-1", "key-1")
|
||||
mock_cost.assert_not_called() # tokens_used=0
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_on_request_success_records_cost_when_configured() -> None:
|
||||
pool = PoolManager(
|
||||
"provider-1",
|
||||
PoolConfig(cost_limit_per_key_tokens=50000),
|
||||
)
|
||||
with (
|
||||
patch(
|
||||
"src.services.provider.pool.redis_ops.set_sticky_binding",
|
||||
new_callable=AsyncMock,
|
||||
),
|
||||
patch(
|
||||
"src.services.provider.pool.redis_ops.touch_lru",
|
||||
new_callable=AsyncMock,
|
||||
),
|
||||
patch(
|
||||
"src.services.provider.pool.redis_ops.add_cost_entry",
|
||||
new_callable=AsyncMock,
|
||||
) as mock_cost,
|
||||
):
|
||||
await pool.on_request_success(session_uuid="sess-1", key_id="key-1", tokens_used=500)
|
||||
|
||||
mock_cost.assert_called_once_with("provider-1", "key-1", 500, 18000)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# on_request_error
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_on_request_error_delegates_to_health_policy(
|
||||
pool: PoolManager,
|
||||
) -> None:
|
||||
with patch(
|
||||
"src.services.provider.pool.health_policy.apply_health_policy",
|
||||
new_callable=AsyncMock,
|
||||
) as mock_hp:
|
||||
await pool.on_request_error(
|
||||
key_id="key-1", status_code=429, error_body=None, response_headers=None
|
||||
)
|
||||
|
||||
mock_hp.assert_called_once()
|
||||
call_kwargs = mock_hp.call_args.kwargs
|
||||
assert call_kwargs["key_id"] == "key-1"
|
||||
assert call_kwargs["status_code"] == 429
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# is_key_schedulable
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_is_key_schedulable_returns_false_when_in_cooldown(
|
||||
pool: PoolManager,
|
||||
) -> None:
|
||||
with patch(
|
||||
"src.services.provider.pool.redis_ops.get_cooldown",
|
||||
new_callable=AsyncMock,
|
||||
return_value="rate_limited_429",
|
||||
):
|
||||
ok, reason = await pool.is_key_schedulable("key-1")
|
||||
|
||||
assert ok is False
|
||||
assert "cooldown" in (reason or "")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_is_key_schedulable_returns_true_when_healthy(
|
||||
pool: PoolManager,
|
||||
) -> None:
|
||||
with patch(
|
||||
"src.services.provider.pool.redis_ops.get_cooldown",
|
||||
new_callable=AsyncMock,
|
||||
return_value=None,
|
||||
):
|
||||
ok, reason = await pool.is_key_schedulable("key-1")
|
||||
|
||||
assert ok is True
|
||||
assert reason is None
|
||||
Reference in New Issue
Block a user