mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-02 01:10:23 +08:00
feat(pool,scheduling): 号池调度维度、配额冷却机制与管理后台重构
- 新增 scheduling_dimensions 模块,为每个 Key 计算多维调度状态(手动/冷却/熔断/成本/健康) - 新增 quota_cooldown 模块,统一判定 Key 的有效冷却原因 - Pool 管理后台 API 扩展 Key 详情字段(调度状态/维度/配额/OAuth 信息) - 前端 Pool 管理页面重写,支持调度状态展示、批量清理封禁 Key - Handler 基类增加请求调度元数据采集,stream telemetry 增强 - 请求时间线组件增强,支持 attempted 候选展示 - Kiro OAuth 凭证导入解析改进 - 新增 usage 表 provider_key 索引迁移 - 补充调度维度、配额冷却、候选枚举等单元测试 Closes #197 Co-authored-by: AAEE86 <33052466+AAEE86@users.noreply.github.com>
This commit is contained in:
@@ -0,0 +1,57 @@
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
import pytest
|
||||
|
||||
from src.services.scheduling.aware_scheduler import CacheAwareScheduler
|
||||
|
||||
|
||||
def _mock_key(key_id: str, api_formats: list[str]) -> MagicMock:
|
||||
key = MagicMock()
|
||||
key.id = key_id
|
||||
key.is_active = True
|
||||
key.api_formats = api_formats
|
||||
key.cache_ttl_minutes = 1
|
||||
key.internal_priority = 1
|
||||
return key
|
||||
|
||||
|
||||
def _mock_endpoint(api_format: str) -> MagicMock:
|
||||
endpoint = MagicMock()
|
||||
endpoint.id = f"ep_{api_format.lower().replace(':', '_')}"
|
||||
endpoint.is_active = True
|
||||
endpoint.api_format = api_format
|
||||
endpoint.api_family = api_format.split(":", 1)[0]
|
||||
endpoint.endpoint_kind = api_format.split(":", 1)[1]
|
||||
endpoint.format_acceptance_config = None
|
||||
return endpoint
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_pool_provider_enumerates_all_key_candidates() -> None:
|
||||
scheduler = CacheAwareScheduler()
|
||||
builder = scheduler._candidate_builder
|
||||
builder._check_model_support = AsyncMock(return_value=(True, None, None, {"m"})) # type: ignore[method-assign]
|
||||
builder._check_key_availability = MagicMock(return_value=(True, None, None)) # type: ignore[method-assign]
|
||||
|
||||
provider = MagicMock()
|
||||
provider.id = "p_pool"
|
||||
provider.name = "pool_provider"
|
||||
provider.enable_format_conversion = False
|
||||
provider.config = {"pool_advanced": {}}
|
||||
provider.endpoints = [_mock_endpoint("openai:chat")]
|
||||
provider.api_keys = [
|
||||
_mock_key("k1", ["openai:chat"]),
|
||||
_mock_key("k2", ["openai:chat"]),
|
||||
]
|
||||
|
||||
candidates = await builder._build_candidates(
|
||||
db=MagicMock(),
|
||||
providers=[provider],
|
||||
client_format="openai:chat",
|
||||
model_name="dummy-model",
|
||||
affinity_key="aff-1",
|
||||
global_conversion_enabled=True,
|
||||
)
|
||||
|
||||
assert len(candidates) == 2
|
||||
assert {str(c.key.id) for c in candidates} == {"k1", "k2"}
|
||||
@@ -310,3 +310,77 @@ async def test_submit_with_failover_filters_missing_billing_rule(
|
||||
assert submit.await_count == 1
|
||||
finally:
|
||||
config.billing_require_rule = old
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_submit_with_failover_applies_pool_reorder_before_submit(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
db = MagicMock()
|
||||
svc = TaskService(db)
|
||||
|
||||
monkeypatch.setattr(
|
||||
"src.services.system.config.SystemConfigService.get_config",
|
||||
lambda *_args, **_kwargs: "provider",
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
"src.services.scheduling.aware_scheduler.get_cache_aware_scheduler",
|
||||
AsyncMock(return_value=None),
|
||||
)
|
||||
|
||||
pool_candidate_a = _make_candidate(
|
||||
provider_id="pool-1",
|
||||
provider_name="pool-provider",
|
||||
endpoint_id="ep-1",
|
||||
key_id="k-a",
|
||||
key_name="key-a",
|
||||
)
|
||||
pool_candidate_b = _make_candidate(
|
||||
provider_id="pool-1",
|
||||
provider_name="pool-provider",
|
||||
endpoint_id="ep-1",
|
||||
key_id="k-b",
|
||||
key_name="key-b",
|
||||
)
|
||||
|
||||
fetch_candidates = AsyncMock(
|
||||
return_value=(
|
||||
[pool_candidate_a, pool_candidate_b],
|
||||
"gm1",
|
||||
)
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
"src.services.candidate.resolver.CandidateResolver.fetch_candidates",
|
||||
fetch_candidates,
|
||||
)
|
||||
|
||||
reordered = [pool_candidate_b, pool_candidate_a]
|
||||
apply_pool_reorder = AsyncMock(return_value=(reordered, []))
|
||||
monkeypatch.setattr(svc, "_apply_pool_reorder", apply_pool_reorder)
|
||||
|
||||
submit = AsyncMock(return_value=httpx.Response(200, json={"id": "task-pooled"}))
|
||||
body = {"session_id": "sid-123"}
|
||||
outcome = await svc.submit_with_failover(
|
||||
api_format="openai:video",
|
||||
model_name="sora",
|
||||
affinity_key="a1",
|
||||
user_api_key=MagicMock(),
|
||||
request_id=None,
|
||||
task_type="video",
|
||||
submit_func=submit,
|
||||
extract_external_task_id=lambda payload: payload.get("id"),
|
||||
supported_auth_types={"api_key"},
|
||||
allow_format_conversion=False,
|
||||
max_candidates=10,
|
||||
request_body=body,
|
||||
)
|
||||
|
||||
assert outcome.external_task_id == "task-pooled"
|
||||
assert outcome.candidate.key.id == "k-b"
|
||||
assert submit.await_count == 1
|
||||
fetch_candidates.assert_awaited_once()
|
||||
assert fetch_candidates.await_args.kwargs.get("request_body") == body
|
||||
apply_pool_reorder.assert_awaited_once_with(
|
||||
[pool_candidate_a, pool_candidate_b],
|
||||
request_body=body,
|
||||
)
|
||||
|
||||
@@ -166,6 +166,43 @@ async def test_trace_build_summary_matches() -> None:
|
||||
assert summary["success_key_id"] == "key-3"[:8]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_trace_build_summary_uses_attempted_key_ids_when_provided() -> None:
|
||||
pool = PoolManager("prov-1", PoolConfig())
|
||||
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, None), "key-2": (None, None)},
|
||||
),
|
||||
patch(
|
||||
"src.services.provider.pool.redis_ops.get_lru_scores",
|
||||
new_callable=AsyncMock,
|
||||
return_value={},
|
||||
),
|
||||
):
|
||||
result = await pool.reorder_candidates(None, [c1, c2])
|
||||
|
||||
trace = getattr(result[0], "_pool_scheduling_trace", None)
|
||||
assert trace is not None
|
||||
|
||||
summary = trace.build_summary(
|
||||
success_key_id="key-1",
|
||||
attempted_key_ids={"key-1"},
|
||||
)
|
||||
assert summary["total_keys"] == 2
|
||||
assert summary["attempted"] == 1
|
||||
assert summary["success_key_id"] == "key-1"[:8]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_sticky_trace_info() -> None:
|
||||
pool = PoolManager("prov-1", PoolConfig(sticky_session_ttl_seconds=3600))
|
||||
|
||||
94
tests/services/test_pool_scheduling_dimensions.py
Normal file
94
tests/services/test_pool_scheduling_dimensions.py
Normal file
@@ -0,0 +1,94 @@
|
||||
"""Tests for pool scheduling dimension evaluation."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from src.services.provider.pool.scheduling_dimensions import (
|
||||
PoolSchedulingDimensionResult,
|
||||
PoolSchedulingSnapshot,
|
||||
evaluate_pool_scheduling_dimensions,
|
||||
list_pool_scheduling_dimensions,
|
||||
summarize_pool_scheduling_dimensions,
|
||||
)
|
||||
|
||||
|
||||
def _snapshot(**overrides: object) -> PoolSchedulingSnapshot:
|
||||
base = {
|
||||
"is_active": True,
|
||||
"cooldown_reason": None,
|
||||
"cooldown_ttl_seconds": None,
|
||||
"circuit_breaker_open": False,
|
||||
"cost_window_usage": 1200,
|
||||
"cost_limit": 10000,
|
||||
"cost_soft_threshold_percent": 80,
|
||||
"health_score": 0.95,
|
||||
}
|
||||
base.update(overrides)
|
||||
return PoolSchedulingSnapshot(**base) # type: ignore[arg-type]
|
||||
|
||||
|
||||
def test_default_dimension_registry_contains_core_dimensions() -> None:
|
||||
names = list_pool_scheduling_dimensions()
|
||||
assert names == ("manual", "cooldown", "circuit", "cost", "health")
|
||||
|
||||
|
||||
def test_summary_available_when_all_dimensions_ok() -> None:
|
||||
dimensions = evaluate_pool_scheduling_dimensions(_snapshot())
|
||||
summary = summarize_pool_scheduling_dimensions(dimensions)
|
||||
|
||||
assert summary.status == "available"
|
||||
assert summary.reason == "available"
|
||||
assert summary.candidate_eligible is True
|
||||
assert summary.blocked_count == 0
|
||||
assert summary.degraded_count == 0
|
||||
assert summary.score == 100.0
|
||||
|
||||
|
||||
def test_summary_blocked_when_manual_disabled() -> None:
|
||||
dimensions = evaluate_pool_scheduling_dimensions(_snapshot(is_active=False))
|
||||
summary = summarize_pool_scheduling_dimensions(dimensions)
|
||||
|
||||
assert summary.status == "blocked"
|
||||
assert summary.reason == "manual_disabled"
|
||||
assert summary.candidate_eligible is False
|
||||
assert summary.blocked_count >= 1
|
||||
assert summary.score < 100.0
|
||||
|
||||
|
||||
def test_summary_degraded_when_cost_reaches_soft_threshold() -> None:
|
||||
dimensions = evaluate_pool_scheduling_dimensions(
|
||||
_snapshot(cost_window_usage=8200, cost_limit=10000, cost_soft_threshold_percent=80)
|
||||
)
|
||||
summary = summarize_pool_scheduling_dimensions(dimensions)
|
||||
|
||||
assert summary.status == "degraded"
|
||||
assert summary.reason == "cost_soft"
|
||||
assert summary.candidate_eligible is True
|
||||
assert summary.blocked_count == 0
|
||||
assert summary.degraded_count >= 1
|
||||
|
||||
|
||||
def test_summary_blocked_on_cooldown_even_if_other_dimensions_ok() -> None:
|
||||
dimensions = evaluate_pool_scheduling_dimensions(
|
||||
_snapshot(cooldown_reason="rate_limited_429", cooldown_ttl_seconds=120)
|
||||
)
|
||||
summary = summarize_pool_scheduling_dimensions(dimensions)
|
||||
|
||||
assert summary.status == "blocked"
|
||||
assert summary.reason == "cooldown"
|
||||
assert summary.candidate_eligible is False
|
||||
assert summary.blocked_count >= 1
|
||||
|
||||
|
||||
def test_empty_summary_defaults_to_available() -> None:
|
||||
summary = summarize_pool_scheduling_dimensions([])
|
||||
assert summary.status == "available"
|
||||
assert summary.reason == "available"
|
||||
assert summary.score == 100.0
|
||||
|
||||
|
||||
def test_dimension_result_keeps_degraded_health_details() -> None:
|
||||
dimensions = evaluate_pool_scheduling_dimensions(_snapshot(health_score=0.65))
|
||||
health = next((item for item in dimensions if item.code == "health_degraded"), None)
|
||||
assert isinstance(health, PoolSchedulingDimensionResult)
|
||||
assert health.status == "degraded"
|
||||
assert health.detail == "0.65"
|
||||
59
tests/services/test_provider_keys_quota_cooldown.py
Normal file
59
tests/services/test_provider_keys_quota_cooldown.py
Normal file
@@ -0,0 +1,59 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from types import SimpleNamespace
|
||||
from typing import Any, cast
|
||||
|
||||
from src.core.provider_types import ProviderType
|
||||
from src.services.provider_keys.quota_cooldown import resolve_effective_cooldown_reason
|
||||
|
||||
|
||||
def _key_with_metadata(upstream_metadata: dict[str, Any]) -> Any:
|
||||
return cast(Any, SimpleNamespace(upstream_metadata=upstream_metadata))
|
||||
|
||||
|
||||
def test_resolve_effective_cooldown_reason_prefers_redis_reason() -> None:
|
||||
key = _key_with_metadata({"codex": {"primary_used_percent": 100.0}})
|
||||
|
||||
reason = resolve_effective_cooldown_reason(
|
||||
provider_type=ProviderType.CODEX,
|
||||
key=key,
|
||||
redis_reason="rate_limited_429",
|
||||
)
|
||||
|
||||
assert reason == "rate_limited_429"
|
||||
|
||||
|
||||
def test_resolve_effective_cooldown_reason_fallbacks_to_codex_quota_exhausted() -> None:
|
||||
key = _key_with_metadata({"codex": {"primary_used_percent": 100.0}})
|
||||
|
||||
reason = resolve_effective_cooldown_reason(
|
||||
provider_type=ProviderType.CODEX,
|
||||
key=key,
|
||||
redis_reason=None,
|
||||
)
|
||||
|
||||
assert reason == "quota_exhausted"
|
||||
|
||||
|
||||
def test_resolve_effective_cooldown_reason_fallbacks_to_kiro_quota_exhausted() -> None:
|
||||
key = _key_with_metadata({"kiro": {"remaining": 0}})
|
||||
|
||||
reason = resolve_effective_cooldown_reason(
|
||||
provider_type=ProviderType.KIRO,
|
||||
key=key,
|
||||
redis_reason=None,
|
||||
)
|
||||
|
||||
assert reason == "quota_exhausted"
|
||||
|
||||
|
||||
def test_resolve_effective_cooldown_reason_returns_none_when_not_exhausted() -> None:
|
||||
key = _key_with_metadata({"codex": {"primary_used_percent": 12.0}})
|
||||
|
||||
reason = resolve_effective_cooldown_reason(
|
||||
provider_type=ProviderType.CODEX,
|
||||
key=key,
|
||||
redis_reason=None,
|
||||
)
|
||||
|
||||
assert reason is None
|
||||
70
tests/services/test_provider_oauth_kiro_import_parsing.py
Normal file
70
tests/services/test_provider_oauth_kiro_import_parsing.py
Normal file
@@ -0,0 +1,70 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
|
||||
from src.api.admin.provider_oauth import _parse_kiro_import_input
|
||||
from src.services.provider.adapters.kiro.models.credentials import KiroAuthConfig
|
||||
|
||||
|
||||
def test_parse_kiro_import_input_array_unwraps_nested_auth_config() -> None:
|
||||
raw = json.dumps(
|
||||
[
|
||||
{
|
||||
"name": "acc-1",
|
||||
"auth_config": {
|
||||
"refresh_token": "rt-1",
|
||||
"auth_method": "identity_center",
|
||||
"client_id": "cid-1",
|
||||
"client_secret": "csec-1",
|
||||
},
|
||||
}
|
||||
]
|
||||
)
|
||||
|
||||
parsed = _parse_kiro_import_input(raw)
|
||||
assert len(parsed) == 1
|
||||
assert parsed[0]["refresh_token"] == "rt-1"
|
||||
assert parsed[0]["auth_method"] == "identity_center"
|
||||
assert parsed[0]["client_id"] == "cid-1"
|
||||
|
||||
|
||||
def test_parse_kiro_import_input_single_object_unwraps_auth_config() -> None:
|
||||
raw = json.dumps(
|
||||
{
|
||||
"name": "acc-1",
|
||||
"authConfig": {
|
||||
"refreshToken": "rt-1",
|
||||
"authType": "builder_id",
|
||||
"clientId": "cid-1",
|
||||
"clientSecret": "csec-1",
|
||||
},
|
||||
}
|
||||
)
|
||||
|
||||
parsed = _parse_kiro_import_input(raw)
|
||||
assert len(parsed) == 1
|
||||
assert parsed[0]["refreshToken"] == "rt-1"
|
||||
assert parsed[0]["authType"] == "builder_id"
|
||||
|
||||
|
||||
def test_kiro_auth_config_from_dict_maps_device_alias_to_idc() -> None:
|
||||
cfg = KiroAuthConfig.from_dict(
|
||||
{
|
||||
"refreshToken": "rt-1",
|
||||
"auth_type": "identity_center",
|
||||
"clientId": "cid-1",
|
||||
"clientSecret": "csec-1",
|
||||
}
|
||||
)
|
||||
assert cfg.auth_method == "idc"
|
||||
|
||||
|
||||
def test_kiro_auth_config_validate_requires_idc_client_fields_when_explicit() -> None:
|
||||
is_valid, message = KiroAuthConfig.validate_required_fields(
|
||||
{
|
||||
"refreshToken": "rt-1",
|
||||
"auth_type": "builder_id",
|
||||
}
|
||||
)
|
||||
assert is_valid is False
|
||||
assert "clientId" in message
|
||||
82
tests/unit/test_monitoring_trace_attempted_subset.py
Normal file
82
tests/unit/test_monitoring_trace_attempted_subset.py
Normal file
@@ -0,0 +1,82 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
from datetime import datetime, timezone
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
from src.api.admin.monitoring.trace import AdminGetRequestTraceAdapter
|
||||
from src.services.request.candidate import RequestCandidateService
|
||||
|
||||
|
||||
def _candidate(*, status: str, latency_ms: int | None = None) -> SimpleNamespace:
|
||||
now = datetime.now(timezone.utc)
|
||||
return SimpleNamespace(
|
||||
id=f"cand-{status}",
|
||||
request_id="req-1",
|
||||
candidate_index=0,
|
||||
retry_index=0,
|
||||
provider_id=None,
|
||||
endpoint_id=None,
|
||||
key_id=None,
|
||||
required_capabilities=None,
|
||||
status=status,
|
||||
skip_reason=None,
|
||||
is_cached=False,
|
||||
status_code=None,
|
||||
error_type=None,
|
||||
error_message=None,
|
||||
latency_ms=latency_ms,
|
||||
concurrent_requests=None,
|
||||
extra_data=None,
|
||||
created_at=now,
|
||||
started_at=None,
|
||||
finished_at=None,
|
||||
)
|
||||
|
||||
|
||||
def _context() -> SimpleNamespace:
|
||||
return SimpleNamespace(
|
||||
db=MagicMock(),
|
||||
add_audit_metadata=lambda **_: None,
|
||||
)
|
||||
|
||||
|
||||
def test_trace_prefers_attempted_subset(monkeypatch: object) -> None:
|
||||
candidates = [
|
||||
_candidate(status="available"),
|
||||
_candidate(status="unused"),
|
||||
_candidate(status="failed", latency_ms=123),
|
||||
]
|
||||
monkeypatch.setattr(
|
||||
RequestCandidateService,
|
||||
"get_candidates_by_request_id",
|
||||
lambda _db, _request_id: candidates,
|
||||
)
|
||||
|
||||
adapter = AdminGetRequestTraceAdapter(request_id="req-1")
|
||||
response = asyncio.run(adapter.handle(_context()))
|
||||
|
||||
assert response.total_candidates == 1
|
||||
assert len(response.candidates) == 1
|
||||
assert response.candidates[0].status == "failed"
|
||||
assert response.total_latency_ms == 123
|
||||
|
||||
|
||||
def test_trace_falls_back_to_all_when_no_attempted(monkeypatch: object) -> None:
|
||||
candidates = [
|
||||
_candidate(status="available"),
|
||||
_candidate(status="unused"),
|
||||
]
|
||||
monkeypatch.setattr(
|
||||
RequestCandidateService,
|
||||
"get_candidates_by_request_id",
|
||||
lambda _db, _request_id: candidates,
|
||||
)
|
||||
|
||||
adapter = AdminGetRequestTraceAdapter(request_id="req-1")
|
||||
response = asyncio.run(adapter.handle(_context()))
|
||||
|
||||
assert response.total_candidates == 2
|
||||
assert len(response.candidates) == 2
|
||||
assert {c.status for c in response.candidates} == {"available", "unused"}
|
||||
128
tests/unit/test_pool_management_scheduling_state.py
Normal file
128
tests/unit/test_pool_management_scheduling_state.py
Normal file
@@ -0,0 +1,128 @@
|
||||
"""Tests for pool management scheduling state assembly."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from types import SimpleNamespace
|
||||
|
||||
from src.api.admin.pool.routes import (
|
||||
_build_pool_scheduling_state,
|
||||
_is_known_banned_key,
|
||||
_is_known_banned_reason,
|
||||
)
|
||||
|
||||
|
||||
def test_pool_scheduling_state_manual_disabled_is_blocked() -> None:
|
||||
(
|
||||
status,
|
||||
reason,
|
||||
_label,
|
||||
reasons,
|
||||
score,
|
||||
candidate_eligible,
|
||||
blocked_count,
|
||||
_degraded_count,
|
||||
dimensions,
|
||||
) = _build_pool_scheduling_state(
|
||||
is_active=False,
|
||||
cooldown_reason=None,
|
||||
cooldown_ttl_seconds=None,
|
||||
circuit_breaker_open=False,
|
||||
cost_window_usage=0,
|
||||
cost_limit=None,
|
||||
cost_soft_threshold_percent=80,
|
||||
health_score=1.0,
|
||||
)
|
||||
|
||||
assert status == "blocked"
|
||||
assert reason == "manual_disabled"
|
||||
assert candidate_eligible is False
|
||||
assert blocked_count >= 1
|
||||
assert score < 100
|
||||
assert any(item.code == "manual_disabled" for item in reasons)
|
||||
assert any(item.code == "manual_disabled" for item in dimensions)
|
||||
|
||||
|
||||
def test_pool_scheduling_state_cooldown_detail_is_mapped() -> None:
|
||||
(
|
||||
_status,
|
||||
reason,
|
||||
_label,
|
||||
reasons,
|
||||
_score,
|
||||
_candidate_eligible,
|
||||
_blocked_count,
|
||||
_degraded_count,
|
||||
dimensions,
|
||||
) = _build_pool_scheduling_state(
|
||||
is_active=True,
|
||||
cooldown_reason="rate_limited_429",
|
||||
cooldown_ttl_seconds=180,
|
||||
circuit_breaker_open=False,
|
||||
cost_window_usage=0,
|
||||
cost_limit=None,
|
||||
cost_soft_threshold_percent=80,
|
||||
health_score=1.0,
|
||||
)
|
||||
|
||||
assert reason == "cooldown"
|
||||
cooldown_reason = next(item for item in reasons if item.code == "cooldown")
|
||||
cooldown_dimension = next(item for item in dimensions if item.code == "cooldown")
|
||||
assert cooldown_reason.detail == "429 限流"
|
||||
assert cooldown_dimension.detail == "429 限流"
|
||||
|
||||
|
||||
def test_pool_scheduling_state_cost_soft_is_degraded() -> None:
|
||||
(
|
||||
status,
|
||||
reason,
|
||||
_label,
|
||||
_reasons,
|
||||
_score,
|
||||
candidate_eligible,
|
||||
blocked_count,
|
||||
degraded_count,
|
||||
_dimensions,
|
||||
) = _build_pool_scheduling_state(
|
||||
is_active=True,
|
||||
cooldown_reason=None,
|
||||
cooldown_ttl_seconds=None,
|
||||
circuit_breaker_open=False,
|
||||
cost_window_usage=85,
|
||||
cost_limit=100,
|
||||
cost_soft_threshold_percent=80,
|
||||
health_score=1.0,
|
||||
)
|
||||
|
||||
assert status == "degraded"
|
||||
assert reason == "cost_soft"
|
||||
assert candidate_eligible is True
|
||||
assert blocked_count == 0
|
||||
assert degraded_count >= 1
|
||||
|
||||
|
||||
def test_known_banned_reason_account_block_prefix() -> None:
|
||||
assert _is_known_banned_reason("[ACCOUNT_BLOCK] Google 要求验证账号") is True
|
||||
|
||||
|
||||
def test_known_banned_key_detects_kiro_banned_metadata() -> None:
|
||||
key = SimpleNamespace(
|
||||
upstream_metadata={"kiro": {"is_banned": True}},
|
||||
oauth_invalid_reason=None,
|
||||
)
|
||||
assert _is_known_banned_key(key, "kiro") is True
|
||||
|
||||
|
||||
def test_known_banned_key_detects_reason_keywords() -> None:
|
||||
key = SimpleNamespace(
|
||||
upstream_metadata={},
|
||||
oauth_invalid_reason="AWS account temporarily suspended",
|
||||
)
|
||||
assert _is_known_banned_key(key, "antigravity") is True
|
||||
|
||||
|
||||
def test_known_banned_key_does_not_treat_token_expired_as_banned() -> None:
|
||||
key = SimpleNamespace(
|
||||
upstream_metadata={"kiro": {"is_banned": False}},
|
||||
oauth_invalid_reason="access token expired",
|
||||
)
|
||||
assert _is_known_banned_key(key, "kiro") is False
|
||||
154
tests/unit/test_request_scheduling_metadata.py
Normal file
154
tests/unit/test_request_scheduling_metadata.py
Normal file
@@ -0,0 +1,154 @@
|
||||
from src.api.handlers.base.base_handler import BaseMessageHandler
|
||||
from src.services.candidate.schema import CandidateKey
|
||||
|
||||
|
||||
def _handler() -> BaseMessageHandler:
|
||||
return BaseMessageHandler.__new__(BaseMessageHandler)
|
||||
|
||||
|
||||
def test_scheduling_audit_detects_internal_failover() -> None:
|
||||
handler = _handler()
|
||||
metadata = handler._merge_scheduling_metadata(
|
||||
{},
|
||||
candidate_keys=[
|
||||
{
|
||||
"candidate_index": 0,
|
||||
"retry_index": 0,
|
||||
"provider_id": "p1",
|
||||
"provider_name": "provider-a",
|
||||
"key_id": "k1",
|
||||
"key_name": "account-a",
|
||||
"status": "failed",
|
||||
"status_code": 429,
|
||||
},
|
||||
{
|
||||
"candidate_index": 1,
|
||||
"retry_index": 0,
|
||||
"provider_id": "p1",
|
||||
"provider_name": "provider-a",
|
||||
"key_id": "k2",
|
||||
"key_name": "account-b",
|
||||
"status": "success",
|
||||
"status_code": 200,
|
||||
},
|
||||
],
|
||||
selected_key_id="k2",
|
||||
fallback_from_request=False,
|
||||
)
|
||||
assert metadata is not None
|
||||
audit = metadata.get("scheduling_audit")
|
||||
assert isinstance(audit, dict)
|
||||
assert audit.get("attempted_count") == 2
|
||||
assert audit.get("account_count") == 2
|
||||
assert audit.get("retry_occurred") is True
|
||||
assert audit.get("failover_occurred") is True
|
||||
assert audit.get("selected_key_id") == "k2"
|
||||
|
||||
accounts = audit.get("accounts")
|
||||
assert isinstance(accounts, list)
|
||||
assert any(a.get("key_id") == "k2" and a.get("selected") for a in accounts)
|
||||
|
||||
|
||||
def test_scheduling_audit_distinguishes_retry_from_failover() -> None:
|
||||
handler = _handler()
|
||||
metadata = handler._merge_scheduling_metadata(
|
||||
{},
|
||||
candidate_keys=[
|
||||
{
|
||||
"candidate_index": 0,
|
||||
"retry_index": 0,
|
||||
"provider_id": "p1",
|
||||
"provider_name": "provider-a",
|
||||
"key_id": "k1",
|
||||
"key_name": "account-a",
|
||||
"status": "failed",
|
||||
"status_code": 503,
|
||||
},
|
||||
{
|
||||
"candidate_index": 0,
|
||||
"retry_index": 1,
|
||||
"provider_id": "p1",
|
||||
"provider_name": "provider-a",
|
||||
"key_id": "k1",
|
||||
"key_name": "account-a",
|
||||
"status": "success",
|
||||
"status_code": 200,
|
||||
},
|
||||
],
|
||||
selected_key_id="k1",
|
||||
fallback_from_request=False,
|
||||
)
|
||||
assert metadata is not None
|
||||
audit = metadata.get("scheduling_audit")
|
||||
assert isinstance(audit, dict)
|
||||
assert audit.get("attempted_count") == 2
|
||||
assert audit.get("account_count") == 1
|
||||
assert audit.get("retry_occurred") is True
|
||||
assert audit.get("failover_occurred") is False
|
||||
|
||||
|
||||
def test_scheduling_metadata_supports_candidate_key_dataclass() -> None:
|
||||
handler = _handler()
|
||||
metadata = handler._merge_scheduling_metadata(
|
||||
{},
|
||||
candidate_keys=[
|
||||
CandidateKey(
|
||||
candidate_index=0,
|
||||
retry_index=0,
|
||||
provider_id="p1",
|
||||
provider_name="provider-a",
|
||||
endpoint_id="e1",
|
||||
key_id="k1",
|
||||
key_name="account-a",
|
||||
status="success",
|
||||
status_code=200,
|
||||
)
|
||||
],
|
||||
selected_key_id="k1",
|
||||
fallback_from_request=False,
|
||||
)
|
||||
assert metadata is not None
|
||||
snapshots = metadata.get("candidate_keys")
|
||||
assert isinstance(snapshots, list)
|
||||
assert snapshots[0]["status"] == "success"
|
||||
assert snapshots[0]["key_id"] == "k1"
|
||||
|
||||
|
||||
def test_scheduling_audit_excludes_unused_candidates() -> None:
|
||||
handler = _handler()
|
||||
metadata = handler._merge_scheduling_metadata(
|
||||
{},
|
||||
candidate_keys=[
|
||||
{
|
||||
"candidate_index": 0,
|
||||
"retry_index": 0,
|
||||
"provider_id": "p1",
|
||||
"provider_name": "provider-a",
|
||||
"key_id": "k1",
|
||||
"key_name": "account-a",
|
||||
"status": "success",
|
||||
"status_code": 200,
|
||||
},
|
||||
{
|
||||
"candidate_index": 1,
|
||||
"retry_index": 0,
|
||||
"provider_id": "p1",
|
||||
"provider_name": "provider-a",
|
||||
"key_id": "k2",
|
||||
"key_name": "account-b",
|
||||
"status": "unused",
|
||||
},
|
||||
],
|
||||
selected_key_id="k1",
|
||||
fallback_from_request=False,
|
||||
)
|
||||
|
||||
assert metadata is not None
|
||||
audit = metadata.get("scheduling_audit")
|
||||
assert isinstance(audit, dict)
|
||||
assert audit.get("attempted_count") == 1
|
||||
assert audit.get("account_count") == 1
|
||||
attempts = audit.get("attempts")
|
||||
assert isinstance(attempts, list)
|
||||
assert len(attempts) == 1
|
||||
assert attempts[0].get("key_id") == "k1"
|
||||
Reference in New Issue
Block a user