mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-01 17:00:21 +08:00
refactor: 提取模型可用性查询到独立模块 ModelAvailabilityQuery
Close #101 - 新增 src/services/model/availability.py，集中管理系统级模型可用性查询条件 - 重构 models_service.py，复用 ModelAvailabilityQuery 的查询方法 - 删除 find_model_by_id 的 provider_model_name 回退逻辑 - 支持 model_mappings 正则匹配（使用 check_model_allowed_with_mappings） - 修复 test_auth.py 中缺少的 is_locked mock 属性 - 新增 ModelAvailabilityQuery 单元测试
This commit is contained in:
fawney19
@@ -235,6 +235,7 @@ class TestAPIKeyAuthentication:
|
||||
|
||||
mock_api_key = MagicMock()
|
||||
mock_api_key.is_active = True
|
||||
mock_api_key.is_locked = False
|
||||
mock_api_key.expires_at = None
|
||||
mock_api_key.user = mock_user
|
||||
mock_api_key.balance_used_usd = 0.0
|
||||
@@ -271,6 +272,7 @@ class TestAPIKeyAuthentication:
|
||||
"""测试 API Key 已禁用"""
|
||||
mock_api_key = MagicMock()
|
||||
mock_api_key.is_active = False
|
||||
mock_api_key.is_locked = False
|
||||
|
||||
mock_db = MagicMock()
|
||||
mock_db.query.return_value.options.return_value.filter.return_value.first.return_value = (
|
||||
@@ -286,6 +288,7 @@ class TestAPIKeyAuthentication:
|
||||
"""测试 API Key 已过期"""
|
||||
mock_api_key = MagicMock()
|
||||
mock_api_key.is_active = True
|
||||
mock_api_key.is_locked = False
|
||||
mock_api_key.expires_at = datetime.now(timezone.utc) - timedelta(days=1)
|
||||
|
||||
mock_db = MagicMock()
|
||||
|
||||
728
tests/services/test_model_availability.py
Normal file
728
tests/services/test_model_availability.py
Normal file
@@ -0,0 +1,728 @@
|
||||
# mypy: disable-error-code="arg-type"
|
||||
"""
|
||||
ModelAvailabilityQuery 单元测试
|
||||
|
||||
说明：
|
||||
- 本测试使用静态源码检查 + FakeSession/FakeQuery 风格，避免引入真实 DB 依赖
|
||||
- 项目日志使用 loguru，pytest 的 caplog 默认无法直接捕获，因此用 loguru sink 验证 warning
|
||||
"""
|
||||
|
||||
import inspect
|
||||
from io import StringIO
|
||||
from typing import Any
|
||||
|
||||
from loguru import logger
|
||||
|
||||
from src.services.model.availability import ModelAvailabilityQuery
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# 测试辅助类
|
||||
# ============================================================================
|
||||
|
||||
|
||||
class FakeQuery:
|
||||
"""通用 Fake Query，支持链式调用和自定义返回数据"""
|
||||
|
||||
def __init__(self, data: list[Any]) -> None:
|
||||
self._data = data
|
||||
|
||||
def filter(self, *_args: Any, **_kwargs: Any) -> "FakeQuery":
|
||||
return self
|
||||
|
||||
def join(self, *_args: Any, **_kwargs: Any) -> "FakeQuery":
|
||||
return self
|
||||
|
||||
def all(self) -> list[Any]:
|
||||
return self._data
|
||||
|
||||
|
||||
class FakeSession:
|
||||
"""
|
||||
通用 Fake Session
|
||||
|
||||
注：FakeSession 仅实现了 Session 的 query 方法，用于测试场景。
|
||||
类型检查会报错，但运行时正常工作。
|
||||
"""
|
||||
|
||||
def __init__(self, query_data: list[Any]) -> None:
|
||||
self._query_data = query_data
|
||||
|
||||
def query(self, *_entities: Any) -> FakeQuery:
|
||||
return FakeQuery(self._query_data)
|
||||
|
||||
|
||||
class TestBaseActiveModelsSourceCode:
|
||||
"""
|
||||
静态验证 base_active_models 的核心实现
|
||||
|
||||
通过检查源码确保关键条件不会被遗漏。
|
||||
"""
|
||||
|
||||
def test_uses_inner_join_for_global_model(self) -> None:
|
||||
"""应使用内连接 GlobalModel（排除 global_model_id=NULL）"""
|
||||
source = inspect.getsource(ModelAvailabilityQuery.base_active_models)
|
||||
|
||||
assert "join(Model.global_model)" in source or ".join(Model.global_model)" in source, (
|
||||
"base_active_models 应使用 join(Model.global_model) 内连接"
|
||||
)
|
||||
assert "outerjoin" not in source.lower(), (
|
||||
"base_active_models 不应使用 outerjoin（会返回 global_model_id=NULL 的记录）"
|
||||
)
|
||||
|
||||
def test_filters_is_available_true_or_null(self) -> None:
|
||||
"""应过滤 is_available = True 或 NULL"""
|
||||
source = inspect.getsource(ModelAvailabilityQuery.base_active_models)
|
||||
|
||||
assert "or_(" in source, "base_active_models 应使用 or_() 处理 is_available"
|
||||
assert "is_available" in source, "base_active_models 应包含 is_available 条件"
|
||||
assert "is_(True)" in source and "is_(None)" in source, (
|
||||
"base_active_models 应同时检查 is_available = True 和 NULL"
|
||||
)
|
||||
|
||||
def test_filters_all_is_active_fields(self) -> None:
|
||||
"""应过滤 Model/Provider/GlobalModel 的 is_active"""
|
||||
source = inspect.getsource(ModelAvailabilityQuery.base_active_models)
|
||||
|
||||
assert "Model.is_active" in source, "base_active_models 应检查 Model.is_active"
|
||||
assert "Provider.is_active" in source, "base_active_models 应检查 Provider.is_active"
|
||||
assert "GlobalModel.is_active" in source, "base_active_models 应检查 GlobalModel.is_active"
|
||||
|
||||
|
||||
class TestGetProviderKeyRules:
|
||||
"""测试 get_provider_key_rules"""
|
||||
|
||||
def test_empty_provider_ids_returns_empty_dict(self) -> None:
|
||||
"""空 provider_ids 应返回空字典"""
|
||||
result = ModelAvailabilityQuery.get_provider_key_rules(
|
||||
FakeSession([]),
|
||||
provider_ids=set(),
|
||||
api_formats=["OPENAI"],
|
||||
provider_to_endpoint_formats={},
|
||||
)
|
||||
assert result == {}
|
||||
|
||||
def test_skips_key_with_invalid_allowed_models_type(self) -> None:
|
||||
"""allowed_models 类型异常时应跳过该 Key 并打日志"""
|
||||
log_output = StringIO()
|
||||
handler_id = logger.add(
|
||||
log_output,
|
||||
format="{message}",
|
||||
level="WARNING",
|
||||
filter=lambda record: "[ModelAvailability]" in record["message"],
|
||||
)
|
||||
|
||||
try:
|
||||
# (key_id, provider_id, allowed_models, api_formats)
|
||||
data = [("key-1", "provider-1", "invalid-string-type", ["OPENAI"])]
|
||||
|
||||
result = ModelAvailabilityQuery.get_provider_key_rules(
|
||||
FakeSession(data),
|
||||
provider_ids={"provider-1"},
|
||||
api_formats=["OPENAI"],
|
||||
provider_to_endpoint_formats={"provider-1": {"OPENAI"}},
|
||||
)
|
||||
|
||||
assert "provider-1" not in result or len(result.get("provider-1", [])) == 0
|
||||
|
||||
log_content = log_output.getvalue()
|
||||
assert "allowed_models 类型异常" in log_content
|
||||
assert "key-1" in log_content
|
||||
finally:
|
||||
logger.remove(handler_id)
|
||||
|
||||
def test_skips_key_with_invalid_api_formats_type(self) -> None:
|
||||
"""api_formats 类型异常时应跳过该 Key 并打日志"""
|
||||
log_output = StringIO()
|
||||
handler_id = logger.add(
|
||||
log_output,
|
||||
format="{message}",
|
||||
level="WARNING",
|
||||
filter=lambda record: "[ModelAvailability]" in record["message"],
|
||||
)
|
||||
|
||||
try:
|
||||
# api_formats 是字符串而非列表
|
||||
data = [("key-1", "provider-1", None, "OPENAI")]
|
||||
|
||||
result = ModelAvailabilityQuery.get_provider_key_rules(
|
||||
FakeSession(data),
|
||||
provider_ids={"provider-1"},
|
||||
api_formats=["OPENAI"],
|
||||
provider_to_endpoint_formats={"provider-1": {"OPENAI"}},
|
||||
)
|
||||
|
||||
assert "provider-1" not in result or len(result.get("provider-1", [])) == 0
|
||||
|
||||
log_content = log_output.getvalue()
|
||||
assert "api_formats 类型异常" in log_content
|
||||
finally:
|
||||
logger.remove(handler_id)
|
||||
|
||||
def test_skips_key_with_none_api_formats(self) -> None:
|
||||
"""api_formats 为 None 时应跳过该 Key（不打日志）"""
|
||||
data = [("key-1", "provider-1", None, None)]
|
||||
|
||||
result = ModelAvailabilityQuery.get_provider_key_rules(
|
||||
FakeSession(data),
|
||||
provider_ids={"provider-1"},
|
||||
api_formats=["OPENAI"],
|
||||
provider_to_endpoint_formats={"provider-1": {"OPENAI"}},
|
||||
)
|
||||
|
||||
assert "provider-1" not in result or len(result.get("provider-1", [])) == 0
|
||||
|
||||
def test_allowed_models_none_means_no_restriction(self) -> None:
|
||||
"""allowed_models = None 表示不限制模型"""
|
||||
data = [("key-1", "provider-1", None, ["OPENAI"])]
|
||||
|
||||
result = ModelAvailabilityQuery.get_provider_key_rules(
|
||||
FakeSession(data),
|
||||
provider_ids={"provider-1"},
|
||||
api_formats=["OPENAI"],
|
||||
provider_to_endpoint_formats={"provider-1": {"OPENAI"}},
|
||||
)
|
||||
|
||||
assert "provider-1" in result
|
||||
rules = result["provider-1"]
|
||||
assert len(rules) == 1
|
||||
allowed_models, usable_formats = rules[0]
|
||||
assert allowed_models is None # None = 不限制
|
||||
assert "OPENAI" in usable_formats
|
||||
|
||||
def test_allowed_models_list_is_preserved(self) -> None:
|
||||
"""allowed_models 为有效列表时应正常返回"""
|
||||
data = [("key-1", "provider-1", ["claude-3-opus", "claude-3-sonnet"], ["OPENAI"])]
|
||||
|
||||
result = ModelAvailabilityQuery.get_provider_key_rules(
|
||||
FakeSession(data),
|
||||
provider_ids={"provider-1"},
|
||||
api_formats=["OPENAI"],
|
||||
provider_to_endpoint_formats={"provider-1": {"OPENAI"}},
|
||||
)
|
||||
|
||||
assert "provider-1" in result
|
||||
rules = result["provider-1"]
|
||||
assert len(rules) == 1
|
||||
allowed_models, _ = rules[0]
|
||||
assert allowed_models == ["claude-3-opus", "claude-3-sonnet"]
|
||||
|
||||
def test_skips_key_when_format_intersection_empty(self) -> None:
|
||||
"""格式交集为空时不包含该 Key"""
|
||||
# Key 支持 CLAUDE，但请求的是 OPENAI
|
||||
data = [("key-1", "provider-1", None, ["CLAUDE"])]
|
||||
|
||||
result = ModelAvailabilityQuery.get_provider_key_rules(
|
||||
FakeSession(data),
|
||||
provider_ids={"provider-1"},
|
||||
api_formats=["OPENAI"],
|
||||
provider_to_endpoint_formats={"provider-1": {"OPENAI"}},
|
||||
)
|
||||
|
||||
assert "provider-1" not in result or len(result.get("provider-1", [])) == 0
|
||||
|
||||
def test_multiple_keys_same_provider(self) -> None:
|
||||
"""同一 Provider 下多个 Key 应合并规则"""
|
||||
data = [
|
||||
("key-1", "provider-1", None, ["OPENAI"]),
|
||||
("key-2", "provider-1", ["claude-3-opus"], ["OPENAI"]),
|
||||
]
|
||||
|
||||
result = ModelAvailabilityQuery.get_provider_key_rules(
|
||||
FakeSession(data),
|
||||
provider_ids={"provider-1"},
|
||||
api_formats=["OPENAI"],
|
||||
provider_to_endpoint_formats={"provider-1": {"OPENAI"}},
|
||||
)
|
||||
|
||||
assert "provider-1" in result
|
||||
rules = result["provider-1"]
|
||||
assert len(rules) == 2
|
||||
|
||||
|
||||
class TestGetProvidersWithActiveKeys:
|
||||
"""测试 get_providers_with_active_keys"""
|
||||
|
||||
def test_empty_provider_ids_returns_empty_set(self) -> None:
|
||||
"""空 provider_ids 应返回空集合"""
|
||||
result = ModelAvailabilityQuery.get_providers_with_active_keys(
|
||||
FakeSession([]),
|
||||
provider_ids=set(),
|
||||
api_formats=["OPENAI"],
|
||||
provider_to_endpoint_formats={},
|
||||
)
|
||||
assert result == set()
|
||||
|
||||
def test_skips_key_with_invalid_api_formats_type(self) -> None:
|
||||
"""api_formats 类型异常时应跳过该 Key 并打日志"""
|
||||
log_output = StringIO()
|
||||
handler_id = logger.add(
|
||||
log_output,
|
||||
format="{message}",
|
||||
level="WARNING",
|
||||
filter=lambda record: "[ModelAvailability]" in record["message"],
|
||||
)
|
||||
|
||||
try:
|
||||
# (provider_id, api_formats) - api_formats 是字符串
|
||||
data = [("provider-1", "OPENAI")]
|
||||
|
||||
result = ModelAvailabilityQuery.get_providers_with_active_keys(
|
||||
FakeSession(data),
|
||||
provider_ids={"provider-1"},
|
||||
api_formats=["OPENAI"],
|
||||
provider_to_endpoint_formats={"provider-1": {"OPENAI"}},
|
||||
)
|
||||
|
||||
assert "provider-1" not in result
|
||||
|
||||
log_content = log_output.getvalue()
|
||||
assert "api_formats 类型异常" in log_content
|
||||
finally:
|
||||
logger.remove(handler_id)
|
||||
|
||||
def test_skips_key_with_none_api_formats(self) -> None:
|
||||
"""api_formats 为 None 时应跳过该 Key（不打日志）"""
|
||||
data = [("provider-1", None)]
|
||||
|
||||
result = ModelAvailabilityQuery.get_providers_with_active_keys(
|
||||
FakeSession(data),
|
||||
provider_ids={"provider-1"},
|
||||
api_formats=["OPENAI"],
|
||||
provider_to_endpoint_formats={"provider-1": {"OPENAI"}},
|
||||
)
|
||||
|
||||
assert "provider-1" not in result
|
||||
|
||||
def test_returns_provider_when_format_matches(self) -> None:
|
||||
"""格式匹配时应返回该 Provider"""
|
||||
data = [("provider-1", ["OPENAI", "CLAUDE"])]
|
||||
|
||||
result = ModelAvailabilityQuery.get_providers_with_active_keys(
|
||||
FakeSession(data),
|
||||
provider_ids={"provider-1"},
|
||||
api_formats=["OPENAI"],
|
||||
provider_to_endpoint_formats={"provider-1": {"OPENAI"}},
|
||||
)
|
||||
|
||||
assert "provider-1" in result
|
||||
|
||||
def test_skips_provider_when_format_intersection_empty(self) -> None:
|
||||
"""格式交集为空时不返回该 Provider"""
|
||||
# Key 支持 GEMINI，但请求的是 OPENAI，且端点支持 OPENAI
|
||||
data = [("provider-1", ["GEMINI"])]
|
||||
|
||||
result = ModelAvailabilityQuery.get_providers_with_active_keys(
|
||||
FakeSession(data),
|
||||
provider_ids={"provider-1"},
|
||||
api_formats=["OPENAI"],
|
||||
provider_to_endpoint_formats={"provider-1": {"OPENAI"}},
|
||||
)
|
||||
|
||||
assert "provider-1" not in result
|
||||
|
||||
def test_skips_provider_not_in_endpoint_formats(self) -> None:
|
||||
"""Provider 不在 endpoint_formats 中时应跳过"""
|
||||
data = [("provider-1", ["OPENAI"])]
|
||||
|
||||
result = ModelAvailabilityQuery.get_providers_with_active_keys(
|
||||
FakeSession(data),
|
||||
provider_ids={"provider-1"},
|
||||
api_formats=["OPENAI"],
|
||||
provider_to_endpoint_formats={}, # 空，无匹配端点
|
||||
)
|
||||
|
||||
assert "provider-1" not in result
|
||||
|
||||
def test_case_insensitive_format_matching(self) -> None:
|
||||
"""格式匹配应忽略大小写"""
|
||||
data = [("provider-1", ["openai"])] # 小写
|
||||
|
||||
result = ModelAvailabilityQuery.get_providers_with_active_keys(
|
||||
FakeSession(data),
|
||||
provider_ids={"provider-1"},
|
||||
api_formats=["OPENAI"], # 大写
|
||||
provider_to_endpoint_formats={"provider-1": {"OPENAI"}},
|
||||
)
|
||||
|
||||
assert "provider-1" in result
|
||||
|
||||
|
||||
class TestGetProvidersWithActiveEndpointsSourceCode:
|
||||
"""静态验证 get_providers_with_active_endpoints 的核心实现"""
|
||||
|
||||
def test_checks_provider_is_active(self) -> None:
|
||||
"""应检查 Provider.is_active"""
|
||||
source = inspect.getsource(ModelAvailabilityQuery.get_providers_with_active_endpoints)
|
||||
assert "Provider.is_active" in source
|
||||
|
||||
def test_checks_endpoint_is_active(self) -> None:
|
||||
"""应检查 ProviderEndpoint.is_active"""
|
||||
source = inspect.getsource(ModelAvailabilityQuery.get_providers_with_active_endpoints)
|
||||
assert "ProviderEndpoint.is_active" in source
|
||||
|
||||
def test_joins_provider_table(self) -> None:
|
||||
"""应 JOIN Provider 表"""
|
||||
source = inspect.getsource(ModelAvailabilityQuery.get_providers_with_active_endpoints)
|
||||
assert ".join(" in source.lower() or "join(" in source.lower()
|
||||
|
||||
|
||||
class TestFindByGlobalModelNameSourceCode:
|
||||
"""静态验证 find_by_global_model_name 的核心实现"""
|
||||
|
||||
def test_uses_base_active_models(self) -> None:
|
||||
"""应复用 base_active_models"""
|
||||
source = inspect.getsource(ModelAvailabilityQuery.find_by_global_model_name)
|
||||
assert "base_active_models" in source
|
||||
|
||||
def test_filters_by_global_model_name(self) -> None:
|
||||
"""应按 GlobalModel.name 过滤"""
|
||||
source = inspect.getsource(ModelAvailabilityQuery.find_by_global_model_name)
|
||||
assert "GlobalModel.name" in source
|
||||
|
||||
|
||||
class TestFindModelByIdNoFallback:
|
||||
"""测试 find_model_by_id 不再回退到 provider_model_name"""
|
||||
|
||||
def test_source_code_no_provider_model_name_in_function(self) -> None:
|
||||
"""
|
||||
静态验证：find_model_by_id 函数中不应出现 provider_model_name
|
||||
|
||||
直接检查字符串是否出现，覆盖所有可能的写法。
|
||||
"""
|
||||
from src.api.base import models_service
|
||||
|
||||
source = inspect.getsource(models_service.find_model_by_id)
|
||||
|
||||
assert "provider_model_name" not in source, (
|
||||
"find_model_by_id 中不应出现 provider_model_name（已删除回退逻辑）"
|
||||
)
|
||||
|
||||
|
||||
class TestGetAvailableModelIdsWithMappings:
|
||||
"""测试 _get_available_model_ids_for_format 支持 model_mappings"""
|
||||
|
||||
def test_source_code_uses_check_model_allowed_with_mappings(self) -> None:
|
||||
"""静态验证：应使用 check_model_allowed_with_mappings 而非 check_model_allowed"""
|
||||
from src.api.base import models_service
|
||||
|
||||
source = inspect.getsource(models_service._get_available_model_ids_for_format)
|
||||
|
||||
assert "check_model_allowed_with_mappings" in source, (
|
||||
"_get_available_model_ids_for_format 应使用 check_model_allowed_with_mappings"
|
||||
)
|
||||
|
||||
def test_source_code_extracts_model_mappings_from_config(self) -> None:
|
||||
"""静态验证：应从 global_model.config 提取 model_mappings"""
|
||||
from src.api.base import models_service
|
||||
|
||||
source = inspect.getsource(models_service._get_available_model_ids_for_format)
|
||||
|
||||
assert "model_mappings" in source, (
|
||||
"_get_available_model_ids_for_format 应提取 model_mappings"
|
||||
)
|
||||
assert ".config" in source or "config" in source, (
|
||||
"_get_available_model_ids_for_format 应从 config 获取 model_mappings"
|
||||
)
|
||||
|
||||
|
||||
class TestModelMappingsIntegration:
|
||||
"""
|
||||
测试 model_mappings 正则映射的集成场景
|
||||
|
||||
场景：
|
||||
- GlobalModel.name = "claude-haiku-4-5-20251001"
|
||||
- Model.provider_model_name = "claude-haiku-4-5-20251001"
|
||||
- GlobalModel.config.model_mappings = ["claude-3-5-haiku-.*"]
|
||||
- Key.allowed_models = ["claude-3-5-haiku-20251001"]
|
||||
|
||||
预期：通过 model_mappings 正则匹配，模型应出现在可用列表中
|
||||
"""
|
||||
|
||||
def test_mapping_pattern_matches_allowed_model(self) -> None:
|
||||
"""model_mappings 正则应能匹配 Key.allowed_models 中的模型"""
|
||||
from src.core.model_permissions import check_model_allowed_with_mappings
|
||||
|
||||
# 模拟你描述的场景
|
||||
is_allowed, matched = check_model_allowed_with_mappings(
|
||||
model_name="claude-haiku-4-5-20251001", # GlobalModel.name
|
||||
allowed_models=["claude-3-5-haiku-20251001"], # Key.allowed_models
|
||||
resolved_model_name="claude-haiku-4-5-20251001", # Model.provider_model_name
|
||||
model_mappings=["claude-3-5-haiku-.*"], # GlobalModel.config.model_mappings
|
||||
)
|
||||
|
||||
assert is_allowed is True
|
||||
assert matched == "claude-3-5-haiku-20251001"
|
||||
|
||||
def test_exact_match_takes_priority_over_mapping(self) -> None:
|
||||
"""精确匹配优先于映射匹配"""
|
||||
from src.core.model_permissions import check_model_allowed_with_mappings
|
||||
|
||||
is_allowed, matched = check_model_allowed_with_mappings(
|
||||
model_name="claude-haiku-4-5-20251001",
|
||||
allowed_models=["claude-haiku-4-5-20251001"], # 精确匹配
|
||||
resolved_model_name="claude-haiku-4-5-20251001",
|
||||
model_mappings=["claude-3-5-haiku-.*"],
|
||||
)
|
||||
|
||||
assert is_allowed is True
|
||||
assert matched is None # 精确匹配时 matched 为 None
|
||||
|
||||
def test_no_mapping_match_returns_false(self) -> None:
|
||||
"""映射不匹配时返回 False"""
|
||||
from src.core.model_permissions import check_model_allowed_with_mappings
|
||||
|
||||
is_allowed, matched = check_model_allowed_with_mappings(
|
||||
model_name="claude-haiku-4-5-20251001",
|
||||
allowed_models=["gpt-4o"], # 与映射模式不匹配
|
||||
resolved_model_name="claude-haiku-4-5-20251001",
|
||||
model_mappings=["claude-3-5-haiku-.*"],
|
||||
)
|
||||
|
||||
assert is_allowed is False
|
||||
assert matched is None
|
||||
|
||||
def test_multiple_mappings_first_match_wins(self) -> None:
|
||||
"""多个映射时，按 allowed_models 排序后第一个匹配的生效"""
|
||||
from src.core.model_permissions import check_model_allowed_with_mappings
|
||||
|
||||
is_allowed, matched = check_model_allowed_with_mappings(
|
||||
model_name="target-model",
|
||||
allowed_models=["b-model-1", "a-model-1"],
|
||||
resolved_model_name="target-model",
|
||||
model_mappings=[".*-model-1"], # 匹配两个
|
||||
)
|
||||
|
||||
assert is_allowed is True
|
||||
# 按字母排序，a-model-1 先匹配
|
||||
assert matched == "a-model-1"
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# AccessRestrictions 测试
|
||||
# ============================================================================
|
||||
|
||||
|
||||
class TestAccessRestrictionsFromApiKeyAndUser:
|
||||
"""测试 AccessRestrictions.from_api_key_and_user 合并逻辑"""
|
||||
|
||||
def test_both_none_returns_no_restrictions(self) -> None:
|
||||
"""API Key 和 User 都为 None 时返回无限制"""
|
||||
from src.api.base.models_service import AccessRestrictions
|
||||
|
||||
result = AccessRestrictions.from_api_key_and_user(None, None)
|
||||
|
||||
assert result.allowed_providers is None
|
||||
assert result.allowed_models is None
|
||||
assert result.allowed_api_formats is None
|
||||
|
||||
def test_api_key_restrictions_take_priority(self) -> None:
|
||||
"""API Key 的限制优先于 User 的限制"""
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
from src.api.base.models_service import AccessRestrictions
|
||||
|
||||
api_key = MagicMock()
|
||||
api_key.allowed_providers = ["provider-a"]
|
||||
api_key.allowed_models = ["model-a"]
|
||||
api_key.allowed_api_formats = ["OPENAI"]
|
||||
|
||||
user = MagicMock()
|
||||
user.allowed_providers = ["provider-b"]
|
||||
user.allowed_models = ["model-b"]
|
||||
user.allowed_api_formats = ["CLAUDE"]
|
||||
|
||||
result = AccessRestrictions.from_api_key_and_user(api_key, user)
|
||||
|
||||
assert result.allowed_providers == ["provider-a"]
|
||||
assert result.allowed_models == ["model-a"]
|
||||
assert result.allowed_api_formats == ["OPENAI"]
|
||||
|
||||
def test_user_restrictions_used_when_api_key_has_none(self) -> None:
|
||||
"""API Key 无限制时使用 User 的限制"""
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
from src.api.base.models_service import AccessRestrictions
|
||||
|
||||
api_key = MagicMock()
|
||||
api_key.allowed_providers = None
|
||||
api_key.allowed_models = None
|
||||
api_key.allowed_api_formats = None
|
||||
|
||||
user = MagicMock()
|
||||
user.allowed_providers = ["provider-b"]
|
||||
user.allowed_models = ["model-b"]
|
||||
user.allowed_api_formats = ["CLAUDE"]
|
||||
|
||||
result = AccessRestrictions.from_api_key_and_user(api_key, user)
|
||||
|
||||
assert result.allowed_providers == ["provider-b"]
|
||||
assert result.allowed_models == ["model-b"]
|
||||
assert result.allowed_api_formats == ["CLAUDE"]
|
||||
|
||||
def test_partial_api_key_restrictions(self) -> None:
|
||||
"""API Key 部分限制时，其余字段从 User 获取"""
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
from src.api.base.models_service import AccessRestrictions
|
||||
|
||||
api_key = MagicMock()
|
||||
api_key.allowed_providers = ["provider-a"]
|
||||
api_key.allowed_models = None # 无限制
|
||||
api_key.allowed_api_formats = None # 无限制
|
||||
|
||||
user = MagicMock()
|
||||
user.allowed_providers = ["provider-b"]
|
||||
user.allowed_models = ["model-b"]
|
||||
user.allowed_api_formats = ["CLAUDE"]
|
||||
|
||||
result = AccessRestrictions.from_api_key_and_user(api_key, user)
|
||||
|
||||
assert result.allowed_providers == ["provider-a"] # 来自 API Key
|
||||
assert result.allowed_models == ["model-b"] # 来自 User
|
||||
assert result.allowed_api_formats == ["CLAUDE"] # 来自 User
|
||||
|
||||
def test_only_api_key_provided(self) -> None:
|
||||
"""只提供 API Key 时使用其限制"""
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
from src.api.base.models_service import AccessRestrictions
|
||||
|
||||
api_key = MagicMock()
|
||||
api_key.allowed_providers = ["provider-a"]
|
||||
api_key.allowed_models = ["model-a"]
|
||||
api_key.allowed_api_formats = ["OPENAI"]
|
||||
|
||||
result = AccessRestrictions.from_api_key_and_user(api_key, None)
|
||||
|
||||
assert result.allowed_providers == ["provider-a"]
|
||||
assert result.allowed_models == ["model-a"]
|
||||
assert result.allowed_api_formats == ["OPENAI"]
|
||||
|
||||
def test_only_user_provided(self) -> None:
|
||||
"""只提供 User 时使用其限制"""
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
from src.api.base.models_service import AccessRestrictions
|
||||
|
||||
user = MagicMock()
|
||||
user.allowed_providers = ["provider-b"]
|
||||
user.allowed_models = ["model-b"]
|
||||
user.allowed_api_formats = ["CLAUDE"]
|
||||
|
||||
result = AccessRestrictions.from_api_key_and_user(None, user)
|
||||
|
||||
assert result.allowed_providers == ["provider-b"]
|
||||
assert result.allowed_models == ["model-b"]
|
||||
assert result.allowed_api_formats == ["CLAUDE"]
|
||||
|
||||
|
||||
class TestAccessRestrictionsIsApiFormatAllowed:
|
||||
"""测试 AccessRestrictions.is_api_format_allowed"""
|
||||
|
||||
def test_none_allowed_formats_means_all_allowed(self) -> None:
|
||||
"""allowed_api_formats = None 表示所有格式都允许"""
|
||||
from src.api.base.models_service import AccessRestrictions
|
||||
|
||||
restrictions = AccessRestrictions(allowed_api_formats=None)
|
||||
|
||||
assert restrictions.is_api_format_allowed("OPENAI") is True
|
||||
assert restrictions.is_api_format_allowed("CLAUDE") is True
|
||||
assert restrictions.is_api_format_allowed("GEMINI") is True
|
||||
|
||||
def test_format_in_allowed_list(self) -> None:
|
||||
"""格式在允许列表中时返回 True"""
|
||||
from src.api.base.models_service import AccessRestrictions
|
||||
|
||||
restrictions = AccessRestrictions(allowed_api_formats=["OPENAI", "CLAUDE"])
|
||||
|
||||
assert restrictions.is_api_format_allowed("OPENAI") is True
|
||||
assert restrictions.is_api_format_allowed("CLAUDE") is True
|
||||
|
||||
def test_format_not_in_allowed_list(self) -> None:
|
||||
"""格式不在允许列表中时返回 False"""
|
||||
from src.api.base.models_service import AccessRestrictions
|
||||
|
||||
restrictions = AccessRestrictions(allowed_api_formats=["OPENAI"])
|
||||
|
||||
assert restrictions.is_api_format_allowed("CLAUDE") is False
|
||||
assert restrictions.is_api_format_allowed("GEMINI") is False
|
||||
|
||||
def test_empty_allowed_list_blocks_all(self) -> None:
|
||||
"""空允许列表阻止所有格式"""
|
||||
from src.api.base.models_service import AccessRestrictions
|
||||
|
||||
restrictions = AccessRestrictions(allowed_api_formats=[])
|
||||
|
||||
assert restrictions.is_api_format_allowed("OPENAI") is False
|
||||
assert restrictions.is_api_format_allowed("CLAUDE") is False
|
||||
|
||||
|
||||
class TestAccessRestrictionsIsModelAllowed:
|
||||
"""测试 AccessRestrictions.is_model_allowed"""
|
||||
|
||||
def test_no_restrictions_allows_all(self) -> None:
|
||||
"""无限制时允许所有模型"""
|
||||
from src.api.base.models_service import AccessRestrictions
|
||||
|
||||
restrictions = AccessRestrictions(allowed_providers=None, allowed_models=None)
|
||||
|
||||
assert restrictions.is_model_allowed("claude-3-opus", "provider-a") is True
|
||||
assert restrictions.is_model_allowed("gpt-4", "provider-b") is True
|
||||
|
||||
def test_provider_restriction_blocks_unallowed_provider(self) -> None:
|
||||
"""Provider 限制阻止不在列表中的 Provider"""
|
||||
from src.api.base.models_service import AccessRestrictions
|
||||
|
||||
restrictions = AccessRestrictions(allowed_providers=["provider-a"])
|
||||
|
||||
assert restrictions.is_model_allowed("claude-3-opus", "provider-a") is True
|
||||
assert restrictions.is_model_allowed("claude-3-opus", "provider-b") is False
|
||||
|
||||
def test_model_restriction_blocks_unallowed_model(self) -> None:
|
||||
"""模型限制阻止不在列表中的模型"""
|
||||
from src.api.base.models_service import AccessRestrictions
|
||||
|
||||
restrictions = AccessRestrictions(allowed_models=["claude-3-opus"])
|
||||
|
||||
assert restrictions.is_model_allowed("claude-3-opus", "provider-a") is True
|
||||
assert restrictions.is_model_allowed("gpt-4", "provider-a") is False
|
||||
|
||||
def test_both_restrictions_must_pass(self) -> None:
|
||||
"""Provider 和模型限制都必须通过"""
|
||||
from src.api.base.models_service import AccessRestrictions
|
||||
|
||||
restrictions = AccessRestrictions(
|
||||
allowed_providers=["provider-a"],
|
||||
allowed_models=["claude-3-opus"],
|
||||
)
|
||||
|
||||
# 两者都满足
|
||||
assert restrictions.is_model_allowed("claude-3-opus", "provider-a") is True
|
||||
|
||||
# Provider 不满足
|
||||
assert restrictions.is_model_allowed("claude-3-opus", "provider-b") is False
|
||||
|
||||
# 模型不满足
|
||||
assert restrictions.is_model_allowed("gpt-4", "provider-a") is False
|
||||
|
||||
# 两者都不满足
|
||||
assert restrictions.is_model_allowed("gpt-4", "provider-b") is False
|
||||
|
||||
def test_empty_provider_list_blocks_all(self) -> None:
|
||||
"""空 Provider 列表阻止所有"""
|
||||
from src.api.base.models_service import AccessRestrictions
|
||||
|
||||
restrictions = AccessRestrictions(allowed_providers=[])
|
||||
|
||||
assert restrictions.is_model_allowed("claude-3-opus", "provider-a") is False
|
||||
|
||||
def test_empty_model_list_blocks_all(self) -> None:
|
||||
"""空模型列表阻止所有"""
|
||||
from src.api.base.models_service import AccessRestrictions
|
||||
|
||||
restrictions = AccessRestrictions(allowed_models=[])
|
||||
|
||||
assert restrictions.is_model_allowed("claude-3-opus", "provider-a") is False
|
||||
|
||||
|
||||
Reference in New Issue
Block a user