Gemini MCP Server

"""Tests for CustomProvider functionality.""" import os from unittest.mock import MagicMock, patch import pytest from providers import ModelProviderRegistry from providers.custom import CustomProvider from providers.shared import ProviderType class TestCustomProvider: """Test CustomProvider class functionality.""" def test_provider_initialization_with_params(self): """Test CustomProvider initializes correctly with explicit parameters.""" provider = CustomProvider(api_key="test-key", base_url="http://localhost:11434/v1") assert provider.base_url == "http://localhost:11434/v1" assert provider.api_key == "test-key" assert provider.get_provider_type() == ProviderType.CUSTOM def test_provider_initialization_with_env_vars(self): """Test CustomProvider initializes correctly with environment variables.""" with patch.dict(os.environ, {"CUSTOM_API_URL": "http://localhost:8000/v1", "CUSTOM_API_KEY": "env-key"}): provider = CustomProvider() assert provider.base_url == "http://localhost:8000/v1" assert provider.api_key == "env-key" def test_provider_initialization_missing_url(self): """Test CustomProvider raises error when URL is missing.""" with patch.dict(os.environ, {"CUSTOM_API_URL": ""}, clear=False): with pytest.raises(ValueError, match="Custom API URL must be provided"): CustomProvider(api_key="test-key") def test_validate_model_names_always_true(self): """Test CustomProvider validates model names correctly.""" provider = CustomProvider(api_key="test-key", base_url="http://localhost:11434/v1") # Known model should validate assert provider.validate_model_name("llama3.2") # For custom provider, unknown models return False when not in registry # This is expected behavior - custom models need to be declared in custom_models.json assert not provider.validate_model_name("unknown-model") assert not provider.validate_model_name("anything") def test_get_capabilities_from_registry(self): """Test get_capabilities returns registry capabilities when available.""" # Save original environment original_env = os.environ.get("OPENROUTER_ALLOWED_MODELS") try: # Clear any restrictions os.environ.pop("OPENROUTER_ALLOWED_MODELS", None) provider = CustomProvider(api_key="test-key", base_url="http://localhost:11434/v1") # OpenRouter-backed models should be handled by the OpenRouter provider with pytest.raises(ValueError): provider.get_capabilities("o3") # Test with a custom model from the local registry capabilities = provider.get_capabilities("local-llama") assert capabilities.provider == ProviderType.CUSTOM assert capabilities.context_window > 0 finally: # Restore original environment if original_env is None: os.environ.pop("OPENROUTER_ALLOWED_MODELS", None) else: os.environ["OPENROUTER_ALLOWED_MODELS"] = original_env def test_get_capabilities_generic_fallback(self): """Test get_capabilities raises error for unknown models not in registry.""" provider = CustomProvider(api_key="test-key", base_url="http://localhost:11434/v1") # Unknown models should raise ValueError when not in registry with pytest.raises(ValueError, match="Unsupported model 'unknown-model-xyz' for provider custom"): provider.get_capabilities("unknown-model-xyz") def test_model_alias_resolution(self): """Test model alias resolution works correctly.""" provider = CustomProvider(api_key="test-key", base_url="http://localhost:11434/v1") # Test that aliases resolve properly # "llama" now resolves to "meta-llama/llama-3-70b" (the OpenRouter model) resolved = provider._resolve_model_name("llama") assert resolved == "meta-llama/llama-3-70b" # Test local model alias resolved_local = provider._resolve_model_name("local-llama") assert resolved_local == "llama3.2" def test_no_thinking_mode_support(self): """Custom provider generic capabilities default to no thinking mode.""" provider = CustomProvider(api_key="test-key", base_url="http://localhost:11434/v1") # llama3.2 is a known model that should work assert not provider.get_capabilities("llama3.2").supports_extended_thinking # Unknown models should raise error with pytest.raises(ValueError, match="Unsupported model 'any-model' for provider custom"): provider.get_capabilities("any-model") @patch("providers.custom.OpenAICompatibleProvider.generate_content") def test_generate_content_with_alias_resolution(self, mock_generate): """Test generate_content resolves aliases before calling parent.""" mock_response = MagicMock() mock_generate.return_value = mock_response provider = CustomProvider(api_key="test-key", base_url="http://localhost:11434/v1") # Call with an alias result = provider.generate_content( prompt="test prompt", model_name="llama", temperature=0.7, # This is an alias ) # Verify parent method was called with resolved model name mock_generate.assert_called_once() call_args = mock_generate.call_args # The model_name should be either resolved or passed through assert "model_name" in call_args.kwargs assert result == mock_response class TestCustomProviderRegistration: """Test CustomProvider integration with ModelProviderRegistry.""" def setup_method(self): """Clear registry before each test.""" ModelProviderRegistry.clear_cache() ModelProviderRegistry.unregister_provider(ProviderType.CUSTOM) def teardown_method(self): """Clean up after each test.""" ModelProviderRegistry.clear_cache() ModelProviderRegistry.unregister_provider(ProviderType.CUSTOM) def test_custom_provider_factory_registration(self): """Test custom provider can be registered via factory function.""" def custom_provider_factory(api_key=None): return CustomProvider(api_key="test-key", base_url="http://localhost:11434/v1") with patch.dict(os.environ, {"CUSTOM_API_PLACEHOLDER": "configured"}): ModelProviderRegistry.register_provider(ProviderType.CUSTOM, custom_provider_factory) # Verify provider is available available = ModelProviderRegistry.get_available_providers() assert ProviderType.CUSTOM in available # Verify provider can be retrieved provider = ModelProviderRegistry.get_provider(ProviderType.CUSTOM) assert provider is not None assert isinstance(provider, CustomProvider) def test_dual_provider_setup(self): """Test both OpenRouter and Custom providers can coexist.""" from providers.openrouter import OpenRouterProvider # Create factory for custom provider def custom_provider_factory(api_key=None): return CustomProvider(api_key="", base_url="http://localhost:11434/v1") with patch.dict( os.environ, { "OPENROUTER_API_KEY": "test-openrouter-key", "CUSTOM_API_PLACEHOLDER": "configured", "OPENROUTER_ALLOWED_MODELS": "llama,anthropic/claude-opus-4.1", }, clear=True, ): # Register both providers ModelProviderRegistry.register_provider(ProviderType.OPENROUTER, OpenRouterProvider) ModelProviderRegistry.register_provider(ProviderType.CUSTOM, custom_provider_factory) # Verify both are available available = ModelProviderRegistry.get_available_providers() assert ProviderType.OPENROUTER in available assert ProviderType.CUSTOM in available # Verify both can be retrieved openrouter_provider = ModelProviderRegistry.get_provider(ProviderType.OPENROUTER) custom_provider = ModelProviderRegistry.get_provider(ProviderType.CUSTOM) assert openrouter_provider is not None assert custom_provider is not None assert isinstance(custom_provider, CustomProvider) def test_provider_priority_selection(self): """Test provider selection prioritizes correctly.""" from providers.openrouter import OpenRouterProvider def custom_provider_factory(api_key=None): return CustomProvider(api_key="", base_url="http://localhost:11434/v1") with patch.dict( os.environ, { "OPENROUTER_API_KEY": "test-openrouter-key", "CUSTOM_API_PLACEHOLDER": "configured", "OPENROUTER_ALLOWED_MODELS": "", }, clear=True, ): import utils.model_restrictions utils.model_restrictions._restriction_service = None custom_provider = custom_provider_factory() openrouter_provider = OpenRouterProvider(api_key="test-openrouter-key") assert not custom_provider.validate_model_name("llama") assert openrouter_provider.validate_model_name("llama") class TestConfigureProvidersFunction: """Test the configure_providers function in server.py.""" def setup_method(self): """Clear environment and registry before each test.""" # Store the original providers to restore them later registry = ModelProviderRegistry() self._original_providers = registry._providers.copy() ModelProviderRegistry.clear_cache() for provider_type in ProviderType: ModelProviderRegistry.unregister_provider(provider_type) def teardown_method(self): """Clean up after each test.""" # Restore the original providers that were registered in conftest.py registry = ModelProviderRegistry() ModelProviderRegistry.clear_cache() registry._providers.clear() registry._providers.update(self._original_providers) def test_configure_providers_custom_only(self): """Test configure_providers with only custom URL set.""" from server import configure_providers with patch.dict( os.environ, { "CUSTOM_API_URL": "http://localhost:11434/v1", "CUSTOM_API_KEY": "", # Clear other API keys "GEMINI_API_KEY": "", "OPENAI_API_KEY": "", "OPENROUTER_API_KEY": "", }, clear=True, ): configure_providers() # Verify only custom provider is available available = ModelProviderRegistry.get_available_providers() assert ProviderType.CUSTOM in available assert ProviderType.OPENROUTER not in available def test_configure_providers_openrouter_only(self): """Test configure_providers with only OpenRouter key set.""" from server import configure_providers with patch.dict( os.environ, { "OPENROUTER_API_KEY": "test-key", # Clear other API keys "GEMINI_API_KEY": "", "OPENAI_API_KEY": "", "CUSTOM_API_URL": "", }, clear=True, ): configure_providers() # Verify only OpenRouter provider is available available = ModelProviderRegistry.get_available_providers() assert ProviderType.OPENROUTER in available assert ProviderType.CUSTOM not in available def test_configure_providers_dual_setup(self): """Test configure_providers with both OpenRouter and Custom configured.""" from server import configure_providers with patch.dict( os.environ, { "OPENROUTER_API_KEY": "test-openrouter-key", "CUSTOM_API_URL": "http://localhost:11434/v1", "CUSTOM_API_KEY": "", # Clear other API keys "GEMINI_API_KEY": "", "OPENAI_API_KEY": "", }, clear=True, ): configure_providers() # Verify both providers are available available = ModelProviderRegistry.get_available_providers() assert ProviderType.OPENROUTER in available assert ProviderType.CUSTOM in available def test_configure_providers_no_valid_keys(self): """Test configure_providers raises error when no valid API keys.""" from server import configure_providers with patch.dict( os.environ, {"GEMINI_API_KEY": "", "OPENAI_API_KEY": "", "OPENROUTER_API_KEY": "", "CUSTOM_API_URL": ""}, clear=True, ): with pytest.raises(ValueError, match="At least one API configuration is required"): configure_providers()