mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-03 09:50:21 +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
279
tests/services/test_pool_health_policy.py
Normal file
279
tests/services/test_pool_health_policy.py
Normal file
@@ -0,0 +1,279 @@
|
||||
"""Tests for pool_health_policy.py — error code classification."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from unittest.mock import AsyncMock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
from src.services.provider.pool.config import (
|
||||
PoolConfig,
|
||||
UnschedulableRule,
|
||||
)
|
||||
from src.services.provider.pool.health_policy import (
|
||||
_extract_error_message,
|
||||
_parse_retry_after,
|
||||
apply_health_policy,
|
||||
)
|
||||
|
||||
PID = "provider-test"
|
||||
KID = "key-test"
|
||||
|
||||
|
||||
@pytest.fixture()
|
||||
def config() -> PoolConfig:
|
||||
return PoolConfig()
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Helper tests
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_parse_retry_after_none() -> None:
|
||||
assert _parse_retry_after(None) is None
|
||||
assert _parse_retry_after({}) is None
|
||||
|
||||
|
||||
def test_parse_retry_after_valid() -> None:
|
||||
assert _parse_retry_after({"retry-after": "60"}) == 60
|
||||
assert _parse_retry_after({"Retry-After": "120"}) == 120
|
||||
|
||||
|
||||
def test_parse_retry_after_clamped() -> None:
|
||||
assert _parse_retry_after({"retry-after": "0"}) == 1
|
||||
assert _parse_retry_after({"retry-after": "9999"}) == 3600
|
||||
|
||||
|
||||
def test_parse_retry_after_invalid() -> None:
|
||||
assert _parse_retry_after({"retry-after": "not-a-number"}) is None
|
||||
|
||||
|
||||
def test_extract_error_message_json_error_object() -> None:
|
||||
body = json.dumps({"error": {"message": "bad request"}})
|
||||
assert _extract_error_message(body) == "bad request"
|
||||
|
||||
|
||||
def test_extract_error_message_json_error_string() -> None:
|
||||
body = json.dumps({"error": "something went wrong"})
|
||||
assert _extract_error_message(body) == "something went wrong"
|
||||
|
||||
|
||||
def test_extract_error_message_plain_text() -> None:
|
||||
assert _extract_error_message("plain text error") == "plain text error"
|
||||
|
||||
|
||||
def test_extract_error_message_empty() -> None:
|
||||
assert _extract_error_message(None) == ""
|
||||
assert _extract_error_message("") == ""
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Status code handling
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_401_invalidates_oauth_cache_and_sets_cooldown(config: PoolConfig) -> None:
|
||||
with (
|
||||
patch(
|
||||
"src.services.provider.pool.redis_ops.invalidate_oauth_token_cache",
|
||||
new_callable=AsyncMock,
|
||||
) as mock_inv,
|
||||
patch(
|
||||
"src.services.provider.pool.redis_ops.set_cooldown",
|
||||
new_callable=AsyncMock,
|
||||
) as mock_cd,
|
||||
):
|
||||
await apply_health_policy(
|
||||
provider_id=PID,
|
||||
key_id=KID,
|
||||
status_code=401,
|
||||
error_body=None,
|
||||
response_headers=None,
|
||||
config=config,
|
||||
)
|
||||
|
||||
mock_inv.assert_called_once_with(KID)
|
||||
mock_cd.assert_called_once_with(PID, KID, "auth_failed_401", ttl=60)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_402_sets_long_cooldown(config: PoolConfig) -> None:
|
||||
with patch(
|
||||
"src.services.provider.pool.redis_ops.set_cooldown",
|
||||
new_callable=AsyncMock,
|
||||
) as mock_cd:
|
||||
await apply_health_policy(
|
||||
provider_id=PID,
|
||||
key_id=KID,
|
||||
status_code=402,
|
||||
error_body=None,
|
||||
response_headers=None,
|
||||
config=config,
|
||||
)
|
||||
|
||||
mock_cd.assert_called_once_with(PID, KID, "payment_required_402", ttl=3600)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_403_sets_long_cooldown(config: PoolConfig) -> None:
|
||||
with patch(
|
||||
"src.services.provider.pool.redis_ops.set_cooldown",
|
||||
new_callable=AsyncMock,
|
||||
) as mock_cd:
|
||||
await apply_health_policy(
|
||||
provider_id=PID,
|
||||
key_id=KID,
|
||||
status_code=403,
|
||||
error_body=None,
|
||||
response_headers=None,
|
||||
config=config,
|
||||
)
|
||||
|
||||
mock_cd.assert_called_once_with(PID, KID, "forbidden_403", ttl=3600)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_400_with_org_disabled_pattern(config: PoolConfig) -> None:
|
||||
body = json.dumps({"error": {"message": "Your organization has been disabled"}})
|
||||
with patch(
|
||||
"src.services.provider.pool.redis_ops.set_cooldown",
|
||||
new_callable=AsyncMock,
|
||||
) as mock_cd:
|
||||
await apply_health_policy(
|
||||
provider_id=PID,
|
||||
key_id=KID,
|
||||
status_code=400,
|
||||
error_body=body,
|
||||
response_headers=None,
|
||||
config=config,
|
||||
)
|
||||
|
||||
mock_cd.assert_called_once()
|
||||
assert "account_disabled_400" in mock_cd.call_args.args[2]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_400_without_pattern_does_nothing(config: PoolConfig) -> None:
|
||||
body = json.dumps({"error": {"message": "invalid json field"}})
|
||||
with patch(
|
||||
"src.services.provider.pool.redis_ops.set_cooldown",
|
||||
new_callable=AsyncMock,
|
||||
) as mock_cd:
|
||||
await apply_health_policy(
|
||||
provider_id=PID,
|
||||
key_id=KID,
|
||||
status_code=400,
|
||||
error_body=body,
|
||||
response_headers=None,
|
||||
config=config,
|
||||
)
|
||||
|
||||
mock_cd.assert_not_called()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_429_uses_retry_after_header(config: PoolConfig) -> None:
|
||||
with patch(
|
||||
"src.services.provider.pool.redis_ops.set_cooldown",
|
||||
new_callable=AsyncMock,
|
||||
) as mock_cd:
|
||||
await apply_health_policy(
|
||||
provider_id=PID,
|
||||
key_id=KID,
|
||||
status_code=429,
|
||||
error_body=None,
|
||||
response_headers={"retry-after": "120"},
|
||||
config=config,
|
||||
)
|
||||
|
||||
mock_cd.assert_called_once_with(PID, KID, "rate_limited_429", ttl=120)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_429_falls_back_to_config_default(config: PoolConfig) -> None:
|
||||
with patch(
|
||||
"src.services.provider.pool.redis_ops.set_cooldown",
|
||||
new_callable=AsyncMock,
|
||||
) as mock_cd:
|
||||
await apply_health_policy(
|
||||
provider_id=PID,
|
||||
key_id=KID,
|
||||
status_code=429,
|
||||
error_body=None,
|
||||
response_headers=None,
|
||||
config=config,
|
||||
)
|
||||
|
||||
mock_cd.assert_called_once_with(PID, KID, "rate_limited_429", ttl=300)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_529_uses_overload_cooldown(config: PoolConfig) -> None:
|
||||
with patch(
|
||||
"src.services.provider.pool.redis_ops.set_cooldown",
|
||||
new_callable=AsyncMock,
|
||||
) as mock_cd:
|
||||
await apply_health_policy(
|
||||
provider_id=PID,
|
||||
key_id=KID,
|
||||
status_code=529,
|
||||
error_body=None,
|
||||
response_headers=None,
|
||||
config=config,
|
||||
)
|
||||
|
||||
mock_cd.assert_called_once_with(PID, KID, "overloaded_529", ttl=30)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Unschedulable keyword rules
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_keyword_rule_matches_and_sets_cooldown() -> None:
|
||||
cfg = PoolConfig(
|
||||
unschedulable_rules=[UnschedulableRule(keyword="capacity", duration_minutes=10)]
|
||||
)
|
||||
body = json.dumps({"error": {"message": "Server at capacity, try later"}})
|
||||
with patch(
|
||||
"src.services.provider.pool.redis_ops.set_cooldown",
|
||||
new_callable=AsyncMock,
|
||||
) as mock_cd:
|
||||
await apply_health_policy(
|
||||
provider_id=PID,
|
||||
key_id=KID,
|
||||
status_code=500,
|
||||
error_body=body,
|
||||
response_headers=None,
|
||||
config=cfg,
|
||||
)
|
||||
|
||||
mock_cd.assert_called_once_with(PID, KID, "rule:capacity", ttl=600)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# health_policy_enabled = False
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_disabled_health_policy_does_nothing() -> None:
|
||||
cfg = PoolConfig(health_policy_enabled=False)
|
||||
with patch(
|
||||
"src.services.provider.pool.redis_ops.set_cooldown",
|
||||
new_callable=AsyncMock,
|
||||
) as mock_cd:
|
||||
await apply_health_policy(
|
||||
provider_id=PID,
|
||||
key_id=KID,
|
||||
status_code=429,
|
||||
error_body=None,
|
||||
response_headers=None,
|
||||
config=cfg,
|
||||
)
|
||||
|
||||
mock_cd.assert_not_called()
|
||||
Reference in New Issue
Block a user