Repository: BeehiveInnovations/zen-mcp-server
File: tests/test_listmodels_restrictions.py (Python)

Code before:
    @patch("utils.model_restrictions.get_restriction_service")
    @patch("providers.openrouter_registry.OpenRouterModelRegistry")
    @patch.object(ModelProviderRegistry, "get_available_models")
    @patch.object(ModelProviderRegistry, "get_provider")
    def test_listmodels_respects_openrouter_restrictions(
        self, mock_get_provider, mock_get_models, mock_registry_class, mock_get_restriction
    ):
        """Test that listmodels only shows allowed OpenRouter models."""
        # Set up mock to return only allowed models when restrictions are respected
        # Include both aliased models and full model names without aliases
        self.mock_openrouter.list_models.return_value = [
            "anthropic/claude-opus-4-20240229",  # Has alias "opus"
            "anthropic/claude-sonnet-4-20240229",  # Has alias "sonnet"
            "deepseek/deepseek-r1-0528:free",  # No alias, full name
            "qwen/qwen3-235b-a22b-04-28:free",  # No alias, full name
        ]

        # Mock registry instance
        mock_registry = MagicMock()
        mock_registry_class.return_value = mock_registry

        # Mock resolve method - return config for aliased models, None for others
        def resolve_side_effect(model_name):
            if "opus" in model_name.lower():
                config = MagicMock()
                config.model_name = "anthropic/claude-opus-4-20240229"
                config.context_window = 200000
                return config
            elif "sonnet" in model_name.lower():
                config = MagicMock()
                config.model_name = "anthropic/claude-sonnet-4-20240229"
                config.context_window = 200000
                return config
            return None  # No config for models without aliases

        mock_registry.resolve.side_effect = resolve_side_effect

        # Mock provider registry
        def get_provider_side_effect(provider_type, force_new=False):
            if provider_type == ProviderType.OPENROUTER:
                return self.mock_openrouter
            elif provider_type == ProviderType.GOOGLE:
                return self.mock_gemini
            return None

        mock_get_provider.side_effect = get_provider_side_effect

        # Mock available models
        mock_get_models.return_value = {
            "gemini-2.5-flash": ProviderType.GOOGLE,
            "gemini-2.5-pro": ProviderType.GOOGLE,

Review comment (suggestion):
```suggestion
                config.model_name = "anthropic/claude-opus-4-20250514"
```

Code after:
    @patch("utils.model_restrictions.get_restriction_service")
    @patch("providers.openrouter_registry.OpenRouterModelRegistry")
    @patch.object(ModelProviderRegistry, "get_available_models")
    @patch.object(ModelProviderRegistry, "get_provider")
    def test_listmodels_respects_openrouter_restrictions(
        self, mock_get_provider, mock_get_models, mock_registry_class, mock_get_restriction
    ):
        """Test that listmodels only shows allowed OpenRouter models."""
        # Set up mock to return only allowed models when restrictions are respected
        # Include both aliased models and full model names without aliases
        self.mock_openrouter.list_models.return_value = [
            "anthropic/claude-opus-4",  # Has alias "opus"
            "anthropic/claude-sonnet-4",  # Has alias "sonnet"
            "deepseek/deepseek-r1-0528:free",  # No alias, full name
            "qwen/qwen3-235b-a22b-04-28:free",  # No alias, full name
        ]

        # Mock registry instance
        mock_registry = MagicMock()
        mock_registry_class.return_value = mock_registry

        # Mock resolve method - return config for aliased models, None for others
        def resolve_side_effect(model_name):
            if "opus" in model_name.lower():
                config = MagicMock()
                config.model_name = "anthropic/claude-opus-4-20240229"
                config.context_window = 200000
                return config
            elif "sonnet" in model_name.lower():
                config = MagicMock()
                config.model_name = "anthropic/claude-sonnet-4-20240229"
                config.context_window = 200000
                return config
            return None  # No config for models without aliases

        mock_registry.resolve.side_effect = resolve_side_effect

        # Mock provider registry
        def get_provider_side_effect(provider_type, force_new=False):
            if provider_type == ProviderType.OPENROUTER:
                return self.mock_openrouter
            elif provider_type == ProviderType.GOOGLE:
                return self.mock_gemini
            return None

        mock_get_provider.side_effect = get_provider_side_effect

        # Mock available models
        mock_get_models.return_value = {
            "gemini-2.5-flash": ProviderType.GOOGLE,
            "gemini-2.5-pro": ProviderType.GOOGLE,