mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-02 17:30:23 +08:00
4ea187cfac
- 新增 PoolCandidate 子类，排序阶段作为单候选参与，执行阶段在 pool_keys 内部选择/切换 key - FailoverEngine 新增 _execute_pool_candidate 方法，支持池内 key 级别故障转移与重试 - 提取 _execute_attempt / _attach_attempt_context / _classify_attempt_error 公共方法 - CandidateBuilder 对号池 Provider 构建单个 PoolCandidate（包含所有可用 key） - CandidateSorter 支持 PoolCandidate 独立优先级分组（global_priority / pool_priority） - CandidateResolver PRE_EXPAND 模式按 pool_keys 展开预创建记录，附加 pool_group_id - TaskService._apply_pool_reorder 改为对 PoolCandidate 调用 select_pool_keys - PoolManager 新增 select_pool_keys 方法，复用 reorder_candidates 逻辑 - 新增 global_priority 号池配置字段（前后端同步） - 前端 Timeline 支持按 pool_group_id 分组显示多号池尝试 - 异步提交路径新增 _expand_pool_candidates_for_async_submit 展开逻辑
62 lines
2.0 KiB
Python
62 lines
2.0 KiB
Python
from unittest.mock import AsyncMock, MagicMock
|
|
|
|
import pytest
|
|
|
|
from src.services.scheduling.aware_scheduler import CacheAwareScheduler
|
|
from src.services.scheduling.schemas import PoolCandidate
|
|
|
|
|
|
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_builds_single_pool_candidate() -> 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) == 1
|
|
pool_candidate = candidates[0]
|
|
assert isinstance(pool_candidate, PoolCandidate)
|
|
assert str(pool_candidate.key.id) == "k1"
|
|
assert {str(k.id) for k in pool_candidate.pool_keys} == {"k1", "k2"}
|