mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-02 09:20:22 +08:00
feat: 调度器根据上游配额自动跳过已耗尽的 Key (Kiro/Codex/Antigravity)
在 _check_key_availability 末尾增加 _is_key_quota_exhausted 检查,读取 ProviderAPIKey.upstream_metadata 中各 Provider 的配额信息,配额耗尽时跳过该 Key 并返回可读的跳过原因。同时修正同函数内 logger 调用为 loguru {} 占位符风格。
This commit is contained in:
225
tests/services/test_aware_scheduler_quota_skipping.py
Normal file
225
tests/services/test_aware_scheduler_quota_skipping.py
Normal file
@@ -0,0 +1,225 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
from src.services.cache.aware_scheduler import CacheAwareScheduler
|
||||
|
||||
|
||||
def _make_key(
|
||||
*,
|
||||
upstream_metadata: dict,
|
||||
allowed_models: list[str] | None = None,
|
||||
capabilities: dict[str, bool] | None = None,
|
||||
) -> MagicMock:
|
||||
key = MagicMock()
|
||||
key.id = "k1234567890"
|
||||
key.allowed_models = allowed_models
|
||||
key.capabilities = capabilities or {}
|
||||
key.upstream_metadata = upstream_metadata
|
||||
return key
|
||||
|
||||
|
||||
@patch(
|
||||
"src.services.cache.aware_scheduler.health_monitor.get_circuit_breaker_status",
|
||||
return_value=(True, None),
|
||||
)
|
||||
def test_kiro_quota_remaining_zero_skips(_mock_cb: MagicMock) -> None:
|
||||
scheduler = CacheAwareScheduler()
|
||||
key = _make_key(upstream_metadata={"kiro": {"remaining": 0.0}})
|
||||
|
||||
ok, reason, _mapped = scheduler._check_key_availability(
|
||||
key,
|
||||
api_format="openai:chat",
|
||||
model_name="any-model",
|
||||
provider_type="kiro",
|
||||
)
|
||||
|
||||
assert ok is False
|
||||
assert reason == "Kiro 账号配额剩余 0"
|
||||
|
||||
|
||||
@patch(
|
||||
"src.services.cache.aware_scheduler.health_monitor.get_circuit_breaker_status",
|
||||
return_value=(True, None),
|
||||
)
|
||||
def test_kiro_quota_remaining_positive_allows(_mock_cb: MagicMock) -> None:
|
||||
scheduler = CacheAwareScheduler()
|
||||
key = _make_key(upstream_metadata={"kiro": {"remaining": 1.0}})
|
||||
|
||||
ok, reason, _mapped = scheduler._check_key_availability(
|
||||
key,
|
||||
api_format="openai:chat",
|
||||
model_name="any-model",
|
||||
provider_type="kiro",
|
||||
)
|
||||
|
||||
assert ok is True
|
||||
assert reason is None
|
||||
|
||||
|
||||
@patch(
|
||||
"src.services.cache.aware_scheduler.health_monitor.get_circuit_breaker_status",
|
||||
return_value=(True, None),
|
||||
)
|
||||
def test_codex_weekly_quota_exhausted_skips(_mock_cb: MagicMock) -> None:
|
||||
scheduler = CacheAwareScheduler()
|
||||
key = _make_key(
|
||||
upstream_metadata={
|
||||
"codex": {
|
||||
"primary_used_percent": 100.0,
|
||||
"secondary_used_percent": 10.0,
|
||||
"code_review_used_percent": 100.0,
|
||||
}
|
||||
}
|
||||
)
|
||||
|
||||
ok, reason, _mapped = scheduler._check_key_availability(
|
||||
key,
|
||||
api_format="openai:cli",
|
||||
model_name="any-model",
|
||||
provider_type="codex",
|
||||
)
|
||||
|
||||
assert ok is False
|
||||
assert reason == "Codex 周限额剩余 0%"
|
||||
|
||||
|
||||
@patch(
|
||||
"src.services.cache.aware_scheduler.health_monitor.get_circuit_breaker_status",
|
||||
return_value=(True, None),
|
||||
)
|
||||
def test_codex_5h_quota_exhausted_skips(_mock_cb: MagicMock) -> None:
|
||||
scheduler = CacheAwareScheduler()
|
||||
key = _make_key(
|
||||
upstream_metadata={
|
||||
"codex": {
|
||||
"primary_used_percent": 10.0,
|
||||
"secondary_used_percent": 100.0,
|
||||
}
|
||||
}
|
||||
)
|
||||
|
||||
ok, reason, _mapped = scheduler._check_key_availability(
|
||||
key,
|
||||
api_format="openai:cli",
|
||||
model_name="any-model",
|
||||
provider_type="codex",
|
||||
)
|
||||
|
||||
assert ok is False
|
||||
assert reason == "Codex 5H 限额剩余 0%"
|
||||
|
||||
|
||||
@patch(
|
||||
"src.services.cache.aware_scheduler.health_monitor.get_circuit_breaker_status",
|
||||
return_value=(True, None),
|
||||
)
|
||||
def test_codex_ignores_code_review_quota(_mock_cb: MagicMock) -> None:
|
||||
scheduler = CacheAwareScheduler()
|
||||
key = _make_key(
|
||||
upstream_metadata={
|
||||
"codex": {
|
||||
"primary_used_percent": 10.0,
|
||||
"secondary_used_percent": 20.0,
|
||||
"code_review_used_percent": 100.0,
|
||||
}
|
||||
}
|
||||
)
|
||||
|
||||
ok, reason, _mapped = scheduler._check_key_availability(
|
||||
key,
|
||||
api_format="openai:cli",
|
||||
model_name="any-model",
|
||||
provider_type="codex",
|
||||
)
|
||||
|
||||
assert ok is True
|
||||
assert reason is None
|
||||
|
||||
|
||||
@patch(
|
||||
"src.services.cache.aware_scheduler.health_monitor.get_circuit_breaker_status",
|
||||
return_value=(True, None),
|
||||
)
|
||||
def test_antigravity_model_quota_exhausted_skips(_mock_cb: MagicMock) -> None:
|
||||
scheduler = CacheAwareScheduler()
|
||||
key = _make_key(
|
||||
upstream_metadata={
|
||||
"antigravity": {
|
||||
"quota_by_model": {
|
||||
"ag-model": {"remaining_fraction": 0.0, "used_percent": 100.0},
|
||||
"other": {"remaining_fraction": 1.0, "used_percent": 0.0},
|
||||
}
|
||||
}
|
||||
},
|
||||
)
|
||||
|
||||
ok, reason, _mapped = scheduler._check_key_availability(
|
||||
key,
|
||||
api_format="gemini:chat",
|
||||
model_name="ag-model",
|
||||
provider_type="antigravity",
|
||||
)
|
||||
|
||||
assert ok is False
|
||||
assert reason == "Antigravity 模型 ag-model 配额剩余 0%"
|
||||
|
||||
|
||||
@patch(
|
||||
"src.services.cache.aware_scheduler.health_monitor.get_circuit_breaker_status",
|
||||
return_value=(True, None),
|
||||
)
|
||||
def test_antigravity_other_model_not_exhausted_allows(_mock_cb: MagicMock) -> None:
|
||||
scheduler = CacheAwareScheduler()
|
||||
key = _make_key(
|
||||
upstream_metadata={
|
||||
"antigravity": {
|
||||
"quota_by_model": {
|
||||
"ag-model": {"remaining_fraction": 0.0, "used_percent": 100.0},
|
||||
"other": {"remaining_fraction": 1.0, "used_percent": 0.0},
|
||||
}
|
||||
}
|
||||
},
|
||||
)
|
||||
|
||||
ok, reason, _mapped = scheduler._check_key_availability(
|
||||
key,
|
||||
api_format="gemini:chat",
|
||||
model_name="other",
|
||||
provider_type="antigravity",
|
||||
)
|
||||
|
||||
assert ok is True
|
||||
assert reason is None
|
||||
|
||||
|
||||
@patch(
|
||||
"src.services.cache.aware_scheduler.health_monitor.get_circuit_breaker_status",
|
||||
return_value=(True, None),
|
||||
)
|
||||
def test_antigravity_quota_uses_mapping_matched_model(_mock_cb: MagicMock) -> None:
|
||||
scheduler = CacheAwareScheduler()
|
||||
|
||||
# Request uses GlobalModel.name, but allowed_models only contains provider-side model id.
|
||||
key = _make_key(
|
||||
upstream_metadata={
|
||||
"antigravity": {
|
||||
"quota_by_model": {
|
||||
"ag-model": {"remaining_fraction": 0.0, "used_percent": 100.0},
|
||||
}
|
||||
}
|
||||
},
|
||||
allowed_models=["ag-model"],
|
||||
)
|
||||
|
||||
ok, reason, mapped = scheduler._check_key_availability(
|
||||
key,
|
||||
api_format="gemini:chat",
|
||||
model_name="global-model",
|
||||
model_mappings=["ag-.*"],
|
||||
provider_type="antigravity",
|
||||
)
|
||||
|
||||
assert mapped == "ag-model"
|
||||
assert ok is False
|
||||
assert reason == "Antigravity 模型 ag-model 配额剩余 0%"
|
||||
Reference in New Issue
Block a user