mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-02 17:30:23 +08:00
7b66505634
将新的 (ApiFamily, EndpointKind) / `family:kind` 签名体系应用到整个代码库： - API Handlers: 所有 adapter/handler 使用新的签名格式 - Services: provider, model, usage, cache, auth 等服务层适配 - Database: ProviderEndpoint 新增 api_family/endpoint_kind 字段 - Frontend: Provider 管理、Usage 表格等组件适配 - Tests: 更新所有相关测试用例
728 lines
28 KiB
Python
728 lines
28 KiB
Python
# 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:chat"],
|
||
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:chat"])]
|
||
|
||
result = ModelAvailabilityQuery.get_provider_key_rules(
|
||
FakeSession(data),
|
||
provider_ids={"provider-1"},
|
||
api_formats=["openai:chat"],
|
||
provider_to_endpoint_formats={"provider-1": {"openai:chat"}},
|
||
)
|
||
|
||
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:chat")]
|
||
|
||
result = ModelAvailabilityQuery.get_provider_key_rules(
|
||
FakeSession(data),
|
||
provider_ids={"provider-1"},
|
||
api_formats=["openai:chat"],
|
||
provider_to_endpoint_formats={"provider-1": {"openai:chat"}},
|
||
)
|
||
|
||
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 时应使用 endpoint 支持的格式（不打日志）"""
|
||
data = [("key-1", "provider-1", None, None)]
|
||
|
||
result = ModelAvailabilityQuery.get_provider_key_rules(
|
||
FakeSession(data),
|
||
provider_ids={"provider-1"},
|
||
api_formats=["openai:chat"],
|
||
provider_to_endpoint_formats={"provider-1": {"openai:chat"}},
|
||
)
|
||
|
||
# api_formats=None 表示使用 endpoint_formats
|
||
assert "provider-1" in result
|
||
rules = result["provider-1"]
|
||
assert len(rules) == 1
|
||
allowed_models, usable_formats = rules[0]
|
||
assert "openai:chat" in usable_formats
|
||
|
||
def test_allowed_models_none_means_no_restriction(self) -> None:
|
||
"""allowed_models = None 表示不限制模型"""
|
||
data = [("key-1", "provider-1", None, ["openai:chat"])]
|
||
|
||
result = ModelAvailabilityQuery.get_provider_key_rules(
|
||
FakeSession(data),
|
||
provider_ids={"provider-1"},
|
||
api_formats=["openai:chat"],
|
||
provider_to_endpoint_formats={"provider-1": {"openai:chat"}},
|
||
)
|
||
|
||
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:chat" 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:chat"])]
|
||
|
||
result = ModelAvailabilityQuery.get_provider_key_rules(
|
||
FakeSession(data),
|
||
provider_ids={"provider-1"},
|
||
api_formats=["openai:chat"],
|
||
provider_to_endpoint_formats={"provider-1": {"openai:chat"}},
|
||
)
|
||
|
||
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:chat"])]
|
||
|
||
result = ModelAvailabilityQuery.get_provider_key_rules(
|
||
FakeSession(data),
|
||
provider_ids={"provider-1"},
|
||
api_formats=["openai:chat"],
|
||
provider_to_endpoint_formats={"provider-1": {"openai:chat"}},
|
||
)
|
||
|
||
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:chat"]),
|
||
("key-2", "provider-1", ["claude-3-opus"], ["openai:chat"]),
|
||
]
|
||
|
||
result = ModelAvailabilityQuery.get_provider_key_rules(
|
||
FakeSession(data),
|
||
provider_ids={"provider-1"},
|
||
api_formats=["openai:chat"],
|
||
provider_to_endpoint_formats={"provider-1": {"openai:chat"}},
|
||
)
|
||
|
||
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:chat"],
|
||
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:chat")]
|
||
|
||
result = ModelAvailabilityQuery.get_providers_with_active_keys(
|
||
FakeSession(data),
|
||
provider_ids={"provider-1"},
|
||
api_formats=["openai:chat"],
|
||
provider_to_endpoint_formats={"provider-1": {"openai:chat"}},
|
||
)
|
||
|
||
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 时应使用 endpoint 支持的格式（返回该 Provider）"""
|
||
data = [("provider-1", None)]
|
||
|
||
result = ModelAvailabilityQuery.get_providers_with_active_keys(
|
||
FakeSession(data),
|
||
provider_ids={"provider-1"},
|
||
api_formats=["openai:chat"],
|
||
provider_to_endpoint_formats={"provider-1": {"openai:chat"}},
|
||
)
|
||
|
||
# api_formats=None 表示使用 endpoint_formats，所以应该返回
|
||
assert "provider-1" in result
|
||
|
||
def test_returns_provider_when_format_matches(self) -> None:
|
||
"""格式匹配时应返回该 Provider"""
|
||
data = [("provider-1", ["openai:chat", "claude:chat"])]
|
||
|
||
result = ModelAvailabilityQuery.get_providers_with_active_keys(
|
||
FakeSession(data),
|
||
provider_ids={"provider-1"},
|
||
api_formats=["openai:chat"],
|
||
provider_to_endpoint_formats={"provider-1": {"openai:chat"}},
|
||
)
|
||
|
||
assert "provider-1" in result
|
||
|
||
def test_skips_provider_when_format_intersection_empty(self) -> None:
|
||
"""格式交集为空时不返回该 Provider"""
|
||
# Key 支持 GEMINI，但请求的是 OPENAI，且端点支持 OPENAI
|
||
data = [("provider-1", ["gemini:chat"])]
|
||
|
||
result = ModelAvailabilityQuery.get_providers_with_active_keys(
|
||
FakeSession(data),
|
||
provider_ids={"provider-1"},
|
||
api_formats=["openai:chat"],
|
||
provider_to_endpoint_formats={"provider-1": {"openai:chat"}},
|
||
)
|
||
|
||
assert "provider-1" not in result
|
||
|
||
def test_skips_provider_not_in_endpoint_formats(self) -> None:
|
||
"""Provider 不在 endpoint_formats 中时应跳过"""
|
||
data = [("provider-1", ["openai:chat"])]
|
||
|
||
result = ModelAvailabilityQuery.get_providers_with_active_keys(
|
||
FakeSession(data),
|
||
provider_ids={"provider-1"},
|
||
api_formats=["openai:chat"],
|
||
provider_to_endpoint_formats={}, # 空，无匹配端点
|
||
)
|
||
|
||
assert "provider-1" not in result
|
||
|
||
def test_case_insensitive_format_matching(self) -> None:
|
||
"""格式匹配应忽略大小写"""
|
||
data = [("provider-1", ["openai:chat"])] # 小写
|
||
|
||
result = ModelAvailabilityQuery.get_providers_with_active_keys(
|
||
FakeSession(data),
|
||
provider_ids={"provider-1"},
|
||
api_formats=["OPENAI:CHAT"], # 大写
|
||
provider_to_endpoint_formats={"provider-1": {"openai:chat"}},
|
||
)
|
||
|
||
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
|
||
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"], # 精确匹配
|
||
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"], # 与映射模式不匹配
|
||
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"],
|
||
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:chat"]
|
||
|
||
user = MagicMock()
|
||
user.allowed_providers = ["provider-b"]
|
||
user.allowed_models = ["model-b"]
|
||
user.allowed_api_formats = ["claude:chat"]
|
||
|
||
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:chat"]
|
||
|
||
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:chat"]
|
||
|
||
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:chat"]
|
||
|
||
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:chat"]
|
||
|
||
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:chat"] # 来自 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:chat"]
|
||
|
||
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:chat"]
|
||
|
||
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:chat"]
|
||
|
||
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:chat"]
|
||
|
||
|
||
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:chat") is True
|
||
assert restrictions.is_api_format_allowed("claude:chat") is True
|
||
assert restrictions.is_api_format_allowed("gemini:chat") 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:chat", "claude:chat"])
|
||
|
||
assert restrictions.is_api_format_allowed("openai:chat") is True
|
||
assert restrictions.is_api_format_allowed("claude:chat") 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:chat"])
|
||
|
||
assert restrictions.is_api_format_allowed("claude:chat") is False
|
||
assert restrictions.is_api_format_allowed("gemini:chat") 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:chat") is False
|
||
assert restrictions.is_api_format_allowed("claude:chat") 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
|