Aether/tests/services/test_aware_scheduler_model_support.py

"""
测试 aware_scheduler 的模型支持检查逻辑

场景：
- GlobalModel claude-haiku 配置了 model_mappings: ["haiku", "claude.*haiku.*"]
- Provider A：Model 表有记录关联到 claude-haiku
- Provider B：Model 表没有记录关联到 claude-haiku
- 用户请求 haiku 模型

期望：Provider B 应该被跳过，因为它不支持 haiku 模型
"""

from unittest.mock import MagicMock, patch

import pytest

from src.models.database import GlobalModel, Model, Provider
from src.services.scheduling.aware_scheduler import CacheAwareScheduler


class TestCheckModelSupportForGlobalModel:
    """测试 _check_model_support_for_global_model 方法"""

    @pytest.mark.asyncio
    async def test_provider_without_model_should_return_false(self) -> None:
        """Provider 没有配置对应的 Model 时应该返回 False"""
        scheduler = CacheAwareScheduler()

        # 创建 GlobalModel
        global_model = MagicMock(spec=GlobalModel)
        global_model.id = "gm-haiku"
        global_model.name = "claude-haiku"
        global_model.supported_capabilities = []

        # 创建 Provider（没有任何 Model）
        provider = MagicMock(spec=Provider)
        provider.name = "Provider-B"
        provider.models = []  # 空列表，没有任何模型

        # 创建 mock db session
        db = MagicMock()
        mock_inspect = MagicMock()
        mock_inspect.transient = False
        mock_inspect.detached = False

        with patch("sqlalchemy.inspect", return_value=mock_inspect):
            is_supported, skip_reason, caps, provider_model_names = (
                await scheduler._candidate_builder._check_model_support_for_global_model(
                    db=db,
                    provider=provider,
                    global_model=global_model,
                    model_name="haiku",
                )
            )

        assert is_supported is False
        assert skip_reason == "Provider 未实现此模型"
        assert caps is None
        assert provider_model_names is None

    @pytest.mark.asyncio
    async def test_provider_with_different_model_should_return_false(self) -> None:
        """Provider 配置了其他模型但没有目标模型时应该返回 False"""
        scheduler = CacheAwareScheduler()

        # 创建 GlobalModel
        global_model = MagicMock(spec=GlobalModel)
        global_model.id = "gm-haiku"
        global_model.name = "claude-haiku"
        global_model.supported_capabilities = []

        # 创建另一个 GlobalModel
        other_global_model = MagicMock(spec=GlobalModel)
        other_global_model.id = "gm-sonnet"
        other_global_model.name = "claude-sonnet"

        # 创建 Model（关联到其他 GlobalModel）
        other_model = MagicMock(spec=Model)
        other_model.global_model_id = "gm-sonnet"  # 关联到 sonnet，不是 haiku
        other_model.is_active = True
        other_model.provider_model_name = "claude-3-sonnet-20240229"

        # 创建 Provider（只有 sonnet 模型，没有 haiku）
        provider = MagicMock(spec=Provider)
        provider.name = "Provider-B"
        provider.models = [other_model]

        # 创建 mock db session
        db = MagicMock()
        mock_inspect = MagicMock()
        mock_inspect.transient = False
        mock_inspect.detached = False

        with patch("sqlalchemy.inspect", return_value=mock_inspect):
            is_supported, skip_reason, caps, provider_model_names = (
                await scheduler._candidate_builder._check_model_support_for_global_model(
                    db=db,
                    provider=provider,
                    global_model=global_model,
                    model_name="haiku",
                )
            )

        assert is_supported is False
        assert skip_reason == "Provider 未实现此模型"

    @pytest.mark.asyncio
    async def test_provider_with_matching_model_should_return_true(self) -> None:
        """Provider 配置了目标模型时应该返回 True"""
        scheduler = CacheAwareScheduler()

        # 创建 GlobalModel
        global_model = MagicMock(spec=GlobalModel)
        global_model.id = "gm-haiku"
        global_model.name = "claude-haiku"
        global_model.supported_capabilities = ["cache_1h"]

        # 创建 Model（关联到目标 GlobalModel）
        model = MagicMock(spec=Model)
        model.global_model_id = "gm-haiku"
        model.is_active = True
        model.provider_model_name = "claude-3-haiku-20240307"
        model.provider_model_mappings = None
        model.get_effective_supports_streaming = MagicMock(return_value=True)

        # 创建 Provider
        provider = MagicMock(spec=Provider)
        provider.name = "Provider-A"
        provider.models = [model]

        # 创建 mock db session
        db = MagicMock()
        mock_inspect = MagicMock()
        mock_inspect.transient = False
        mock_inspect.detached = False

        with patch("sqlalchemy.inspect", return_value=mock_inspect):
            is_supported, skip_reason, caps, provider_model_names = (
                await scheduler._candidate_builder._check_model_support_for_global_model(
                    db=db,
                    provider=provider,
                    global_model=global_model,
                    model_name="haiku",
                )
            )

        assert is_supported is True
        assert skip_reason is None
        assert caps == ["cache_1h"]
        assert provider_model_names == {"claude-3-haiku-20240307"}

    @pytest.mark.asyncio
    async def test_provider_with_inactive_model_should_return_false(self) -> None:
        """Provider 的模型未激活时应该返回 False"""
        scheduler = CacheAwareScheduler()

        # 创建 GlobalModel
        global_model = MagicMock(spec=GlobalModel)
        global_model.id = "gm-haiku"
        global_model.name = "claude-haiku"
        global_model.supported_capabilities = []

        # 创建 Model（未激活）
        model = MagicMock(spec=Model)
        model.global_model_id = "gm-haiku"
        model.is_active = False  # 未激活
        model.provider_model_name = "claude-3-haiku-20240307"

        # 创建 Provider
        provider = MagicMock(spec=Provider)
        provider.name = "Provider-A"
        provider.models = [model]

        # 创建 mock db session
        db = MagicMock()
        mock_inspect = MagicMock()
        mock_inspect.transient = False
        mock_inspect.detached = False

        with patch("sqlalchemy.inspect", return_value=mock_inspect):
            is_supported, skip_reason, caps, provider_model_names = (
                await scheduler._candidate_builder._check_model_support_for_global_model(
                    db=db,
                    provider=provider,
                    global_model=global_model,
                    model_name="haiku",
                )
            )

        assert is_supported is False
        assert skip_reason == "Provider 未实现此模型"