mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-02 17:30:23 +08:00
63870931af
- 将调度器相关模块从 src/services/cache/ 迁移到 src/services/scheduling/ - 下沉类型定义到 core 层: AccessRestrictions, ProviderAuthInfo, ParsedChunk/StreamStats, 视频工具函数 - 提取 thinking_cache 签名缓存到 core/api_format/conversion/ - 提取 provider 认证逻辑到 services/provider/auth - 提取遥测记录到 services/usage/telemetry - 提取 models 列表缓存到 services/cache/model_list_cache - 更新所有引用方的 import 路径及相关测试
330 lines
11 KiB
Python
330 lines
11 KiB
Python
from typing import Any
|
||
from unittest.mock import AsyncMock, MagicMock
|
||
|
||
import pytest
|
||
|
||
from src.core.api_format.conversion import register_default_normalizers
|
||
from src.services.scheduling.aware_scheduler import (
|
||
CacheAwareScheduler,
|
||
_sort_endpoints_by_family_priority,
|
||
)
|
||
|
||
|
||
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, config: dict | None = None) -> 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 = config
|
||
return endpoint
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_build_candidates_allows_cross_format_when_endpoint_accepts_and_overrides_off() -> (
|
||
None
|
||
):
|
||
register_default_normalizers()
|
||
|
||
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.name = "p1"
|
||
provider.enable_format_conversion = False
|
||
provider.endpoints = [
|
||
_mock_endpoint(
|
||
"openai:chat",
|
||
{"enabled": True, "accept_formats": ["claude:chat"], "stream_conversion": True},
|
||
)
|
||
]
|
||
provider.api_keys = [_mock_key("k1", ["openai:chat"])]
|
||
|
||
candidates = await builder._build_candidates(
|
||
db=MagicMock(),
|
||
providers=[provider],
|
||
client_format="claude:chat",
|
||
model_name="dummy-model",
|
||
affinity_key=None,
|
||
global_conversion_enabled=True, # 全局开关开启
|
||
)
|
||
|
||
assert len(candidates) == 1
|
||
assert candidates[0].needs_conversion is True
|
||
assert candidates[0].provider_api_format == "openai:chat"
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_build_candidates_allows_cross_format_when_global_off_but_endpoint_enabled() -> None:
|
||
"""
|
||
分层开关设计：全局 OFF 时回退到端点配置
|
||
- 全局 OFF + 端点 enabled=True -> 允许（端点覆盖全局默认）
|
||
"""
|
||
register_default_normalizers()
|
||
|
||
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.name = "p1"
|
||
provider.enable_format_conversion = False
|
||
provider.endpoints = [
|
||
_mock_endpoint(
|
||
"openai:chat",
|
||
{"enabled": True, "accept_formats": ["claude:chat"], "stream_conversion": True},
|
||
)
|
||
]
|
||
provider.api_keys = [_mock_key("k1", ["openai:chat"])]
|
||
|
||
candidates = await builder._build_candidates(
|
||
db=MagicMock(),
|
||
providers=[provider],
|
||
client_format="claude:chat",
|
||
model_name="dummy-model",
|
||
affinity_key=None,
|
||
global_conversion_enabled=False, # 全局开关关闭，但端点配置允许
|
||
)
|
||
|
||
# 新设计：全局 OFF 时回退到端点配置，端点 enabled=True 则允许
|
||
assert len(candidates) == 1
|
||
assert candidates[0].needs_conversion is True
|
||
assert candidates[0].provider_api_format == "openai:chat"
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_build_candidates_blocks_cross_format_when_global_off_and_endpoint_not_configured() -> (
|
||
None
|
||
):
|
||
"""
|
||
分层开关设计：全局 OFF + 端点未配置 -> 阻止
|
||
"""
|
||
register_default_normalizers()
|
||
|
||
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.name = "p1"
|
||
provider.enable_format_conversion = False
|
||
provider.endpoints = [_mock_endpoint("openai:chat", None)] # 端点未配置格式接受策略
|
||
provider.api_keys = [_mock_key("k1", ["openai:chat"])]
|
||
|
||
candidates = await builder._build_candidates(
|
||
db=MagicMock(),
|
||
providers=[provider],
|
||
client_format="claude:chat",
|
||
model_name="dummy-model",
|
||
affinity_key=None,
|
||
global_conversion_enabled=False, # 全局开关关闭
|
||
)
|
||
|
||
# 全局 OFF + 端点未配置 -> 阻止
|
||
assert candidates == []
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_build_candidates_includes_cross_format_when_enabled() -> None:
|
||
register_default_normalizers()
|
||
|
||
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.name = "p1"
|
||
provider.enable_format_conversion = False
|
||
provider.endpoints = [
|
||
# 端点未配置/未启用格式接受策略，但 DB 全局覆盖开启应强制允许
|
||
_mock_endpoint("openai:chat", None)
|
||
]
|
||
provider.api_keys = [_mock_key("k1", ["openai:chat"])]
|
||
|
||
candidates = await builder._build_candidates(
|
||
db=MagicMock(),
|
||
providers=[provider],
|
||
client_format="claude:chat",
|
||
model_name="dummy-model",
|
||
affinity_key=None,
|
||
global_conversion_enabled=True, # 全局开关开启：跳过端点检查
|
||
)
|
||
|
||
assert len(candidates) == 1
|
||
assert candidates[0].needs_conversion is True
|
||
assert candidates[0].provider_api_format == "openai:chat"
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_exact_matches_rank_before_convertible() -> None:
|
||
register_default_normalizers()
|
||
|
||
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.name = "p1"
|
||
provider.enable_format_conversion = False
|
||
# 故意把 OPENAI 放在 endpoints[0]，验证排序仍然是 CLAUDE（exact）在前
|
||
provider.endpoints = [
|
||
_mock_endpoint(
|
||
"openai:chat",
|
||
{"enabled": True, "accept_formats": ["claude:chat"], "stream_conversion": True},
|
||
),
|
||
_mock_endpoint("claude:chat", None),
|
||
]
|
||
provider.api_keys = [
|
||
_mock_key("k_openai", ["openai:chat"]),
|
||
_mock_key("k_claude", ["claude:chat"]),
|
||
]
|
||
|
||
candidates = await builder._build_candidates(
|
||
db=MagicMock(),
|
||
providers=[provider],
|
||
client_format="claude:chat",
|
||
model_name="dummy-model",
|
||
affinity_key=None,
|
||
global_conversion_enabled=True, # 全局开关开启
|
||
)
|
||
|
||
assert len(candidates) == 2
|
||
assert candidates[0].needs_conversion is False
|
||
assert candidates[0].provider_api_format == "claude:chat"
|
||
assert candidates[1].needs_conversion is True
|
||
assert candidates[1].provider_api_format == "openai:chat"
|
||
|
||
|
||
def test_sort_endpoints_by_family_priority_orders_openai_claude_gemini() -> None:
|
||
eps = [
|
||
_mock_endpoint("gemini:chat"),
|
||
_mock_endpoint("claude:chat"),
|
||
_mock_endpoint("openai:chat"),
|
||
]
|
||
result = _sort_endpoints_by_family_priority(eps)
|
||
assert [e.api_family for e in result] == ["openai", "claude", "gemini"]
|
||
|
||
|
||
def test_sort_endpoints_by_family_priority_unknown_family_sorted_last() -> None:
|
||
eps = [
|
||
_mock_endpoint("unknown:chat"),
|
||
_mock_endpoint("openai:chat"),
|
||
]
|
||
result = _sort_endpoints_by_family_priority(eps)
|
||
assert [e.api_family for e in result] == ["openai", "unknown"]
|
||
|
||
|
||
def test_sort_endpoints_by_family_priority_stable_for_same_family() -> None:
|
||
ep1 = _mock_endpoint("openai:chat")
|
||
ep1.base_url = "url_1"
|
||
ep2 = _mock_endpoint("openai:chat")
|
||
ep2.base_url = "url_2"
|
||
result = _sort_endpoints_by_family_priority([ep1, ep2])
|
||
assert result[0].base_url == "url_1"
|
||
assert result[1].base_url == "url_2"
|
||
|
||
|
||
def test_sort_endpoints_by_family_priority_empty_list() -> None:
|
||
assert _sort_endpoints_by_family_priority([]) == []
|
||
|
||
|
||
def _group_and_sort_endpoints(
|
||
client_family: str, client_kind: str, endpoints: list[MagicMock]
|
||
) -> list[Any]:
|
||
preferred, preferred_other, fallback, fallback_other = [], [], [], []
|
||
for ep in endpoints:
|
||
same_family = ep.api_family == client_family
|
||
same_kind = ep.endpoint_kind == client_kind
|
||
if same_family and same_kind:
|
||
preferred.append(ep)
|
||
elif same_kind:
|
||
preferred_other.append(ep)
|
||
elif same_family:
|
||
fallback.append(ep)
|
||
else:
|
||
fallback_other.append(ep)
|
||
|
||
return (
|
||
_sort_endpoints_by_family_priority(preferred)
|
||
+ _sort_endpoints_by_family_priority(preferred_other)
|
||
+ _sort_endpoints_by_family_priority(fallback)
|
||
+ _sort_endpoints_by_family_priority(fallback_other)
|
||
)
|
||
|
||
|
||
def test_group_and_sort_endpoints_client_openai_chat() -> None:
|
||
endpoints = [
|
||
_mock_endpoint("gemini:cli"),
|
||
_mock_endpoint("claude:chat"),
|
||
_mock_endpoint("openai:cli"),
|
||
_mock_endpoint("gemini:chat"),
|
||
_mock_endpoint("claude:cli"),
|
||
_mock_endpoint("openai:chat"),
|
||
]
|
||
result = _group_and_sort_endpoints("openai", "chat", endpoints)
|
||
assert [(e.api_family, e.endpoint_kind) for e in result] == [
|
||
("openai", "chat"),
|
||
("claude", "chat"),
|
||
("gemini", "chat"),
|
||
("openai", "cli"),
|
||
("claude", "cli"),
|
||
("gemini", "cli"),
|
||
]
|
||
|
||
|
||
def test_group_and_sort_endpoints_client_claude_chat() -> None:
|
||
endpoints = [
|
||
_mock_endpoint("gemini:cli"),
|
||
_mock_endpoint("claude:chat"),
|
||
_mock_endpoint("openai:cli"),
|
||
_mock_endpoint("gemini:chat"),
|
||
_mock_endpoint("claude:cli"),
|
||
_mock_endpoint("openai:chat"),
|
||
]
|
||
result = _group_and_sort_endpoints("claude", "chat", endpoints)
|
||
assert [(e.api_family, e.endpoint_kind) for e in result] == [
|
||
("claude", "chat"),
|
||
("openai", "chat"),
|
||
("gemini", "chat"),
|
||
("claude", "cli"),
|
||
("openai", "cli"),
|
||
("gemini", "cli"),
|
||
]
|
||
|
||
|
||
def test_group_and_sort_endpoints_client_openai_cli() -> None:
|
||
endpoints = [
|
||
_mock_endpoint("gemini:cli"),
|
||
_mock_endpoint("claude:chat"),
|
||
_mock_endpoint("openai:cli"),
|
||
_mock_endpoint("gemini:chat"),
|
||
_mock_endpoint("claude:cli"),
|
||
_mock_endpoint("openai:chat"),
|
||
]
|
||
result = _group_and_sort_endpoints("openai", "cli", endpoints)
|
||
assert [(e.api_family, e.endpoint_kind) for e in result] == [
|
||
("openai", "cli"),
|
||
("claude", "cli"),
|
||
("gemini", "cli"),
|
||
("openai", "chat"),
|
||
("claude", "chat"),
|
||
("gemini", "chat"),
|
||
]
|