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

Code before:
        """Clean up after tests."""
        ModelProviderRegistry.clear_cache()
        # Clean up environment variables
        for key in ["OPENROUTER_ALLOWED_MODELS", "OPENROUTER_API_KEY", "GEMINI_API_KEY"]:
            os.environ.pop(key, None)

    @patch.dict(
        os.environ,
        {
            "OPENROUTER_API_KEY": "test-key",
            "OPENROUTER_ALLOWED_MODELS": "opus,sonnet,deepseek/deepseek-r1-0528:free,qwen/qwen3-235b-a22b-04-28:free",
            "GEMINI_API_KEY": "gemini-test-key",
        },
    )
    @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


Review comment (suggestion):
```suggestion
            "anthropic/claude-opus-4-20250514",  # Has alias "opus"
```

Code after:
        """Clean up after tests."""
        ModelProviderRegistry.clear_cache()
        # Clean up environment variables
        for key in ["OPENROUTER_ALLOWED_MODELS", "OPENROUTER_API_KEY", "GEMINI_API_KEY"]:
            os.environ.pop(key, None)

    @patch.dict(
        os.environ,
        {
            "OPENROUTER_API_KEY": "test-key",
            "OPENROUTER_ALLOWED_MODELS": "opus,sonnet,deepseek/deepseek-r1-0528:free,qwen/qwen3-235b-a22b-04-28:free",
            "GEMINI_API_KEY": "gemini-test-key",
        },
    )
    @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