Gemini MCP Server

"""Test listmodels tool respects model restrictions.""" import asyncio import os import unittest from unittest.mock import MagicMock, patch from providers.base import ModelProvider from providers.registry import ModelProviderRegistry from providers.shared import ModelCapabilities, ProviderType from tools.listmodels import ListModelsTool class TestListModelsRestrictions(unittest.TestCase): """Test that listmodels respects OPENROUTER_ALLOWED_MODELS.""" def setUp(self): """Set up test environment.""" # Clear any existing registry state ModelProviderRegistry.clear_cache() # Create mock OpenRouter provider self.mock_openrouter = MagicMock(spec=ModelProvider) self.mock_openrouter.provider_type = ProviderType.OPENROUTER def make_capabilities( canonical: str, friendly: str, *, aliases=None, context: int = 200_000 ) -> ModelCapabilities: return ModelCapabilities( provider=ProviderType.OPENROUTER, model_name=canonical, friendly_name=friendly, intelligence_score=20, description=friendly, aliases=aliases or [], context_window=context, max_output_tokens=context, supports_extended_thinking=True, ) opus_caps = make_capabilities( "anthropic/claude-opus-4-20240229", "Claude Opus", aliases=["opus"], ) sonnet_caps = make_capabilities( "anthropic/claude-sonnet-4-20240229", "Claude Sonnet", aliases=["sonnet"], ) deepseek_caps = make_capabilities( "deepseek/deepseek-r1-0528:free", "DeepSeek R1", aliases=[], ) qwen_caps = make_capabilities( "qwen/qwen3-235b-a22b-04-28:free", "Qwen3", aliases=[], ) self._openrouter_caps_map = { "anthropic/claude-opus-4": opus_caps, "opus": opus_caps, "anthropic/claude-opus-4-20240229": opus_caps, "anthropic/claude-sonnet-4": sonnet_caps, "sonnet": sonnet_caps, "anthropic/claude-sonnet-4-20240229": sonnet_caps, "deepseek/deepseek-r1-0528:free": deepseek_caps, "qwen/qwen3-235b-a22b-04-28:free": qwen_caps, } self.mock_openrouter.get_capabilities.side_effect = self._openrouter_caps_map.__getitem__ self.mock_openrouter.get_capabilities_by_rank.return_value = [] self.mock_openrouter.list_models.return_value = [] # Create mock Gemini provider for comparison self.mock_gemini = MagicMock(spec=ModelProvider) self.mock_gemini.provider_type = ProviderType.GOOGLE self.mock_gemini.list_models.return_value = ["gemini-2.5-flash", "gemini-2.5-pro"] self.mock_gemini.get_capabilities_by_rank.return_value = [] self.mock_gemini.get_capabilities_by_rank.return_value = [] def tearDown(self): """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.registries.openrouter.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 config.get_effective_capability_rank.return_value = 90 # High rank for Opus return config elif "sonnet" in model_name.lower(): config = MagicMock() config.model_name = "anthropic/claude-sonnet-4-20240229" config.context_window = 200000 config.get_effective_capability_rank.return_value = 80 # Lower rank for Sonnet return config elif "deepseek" in model_name.lower(): config = MagicMock() config.model_name = "deepseek/deepseek-r1-0528:free" config.context_window = 100000 config.get_effective_capability_rank.return_value = 70 return config elif "qwen" in model_name.lower(): config = MagicMock() config.model_name = "qwen/qwen3-235b-a22b-04-28:free" config.context_window = 100000 config.get_effective_capability_rank.return_value = 60 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 # Ensure registry is cleared before test ModelProviderRegistry._registry = {} # Mock available models mock_get_models.return_value = { "gemini-2.5-flash": ProviderType.GOOGLE, "gemini-2.5-pro": ProviderType.GOOGLE, "anthropic/claude-opus-4-20240229": ProviderType.OPENROUTER, "anthropic/claude-sonnet-4-20240229": ProviderType.OPENROUTER, "deepseek/deepseek-r1-0528:free": ProviderType.OPENROUTER, "qwen/qwen3-235b-a22b-04-28:free": ProviderType.OPENROUTER, } # Mock restriction service mock_restriction_service = MagicMock() mock_restriction_service.has_restrictions.return_value = True mock_restriction_service.get_allowed_models.return_value = { "opus", "sonnet", "deepseek/deepseek-r1-0528:free", "qwen/qwen3-235b-a22b-04-28:free", } mock_get_restriction.return_value = mock_restriction_service # Create tool and execute tool = ListModelsTool() # Execute asynchronously loop = asyncio.new_event_loop() asyncio.set_event_loop(loop) result_contents = loop.run_until_complete(tool.execute({})) loop.close() # Extract text content from result result_text = result_contents[0].text # Parse JSON response import json result_json = json.loads(result_text) result = result_json["content"] # Parse the output lines = result.split("\n") # Debug: print the actual result for troubleshooting # print(f"DEBUG: Full result:\n{result}") # Check that OpenRouter section exists openrouter_section_found = False openrouter_models = [] in_openrouter_section = False for line in lines: if "OpenRouter" in line and "✅" in line: openrouter_section_found = True elif ("Models (policy restricted)" in line or "Available Models" in line) and openrouter_section_found: in_openrouter_section = True elif in_openrouter_section: # Check for lines with model names in backticks # Format: - `model-name` (score X) if line.strip().startswith("- ") and "`" in line: # Extract model name between backticks parts = line.split("`") if len(parts) >= 2: model_name = parts[1] openrouter_models.append(model_name) # Stop parsing when we hit the next section elif "##" in line and in_openrouter_section: break self.assertTrue(openrouter_section_found, "OpenRouter section not found") self.assertEqual( len(openrouter_models), 4, f"Expected 4 models, got {len(openrouter_models)}: {openrouter_models}" ) # Verify we did not fall back to unrestricted listing self.mock_openrouter.list_models.assert_not_called() # Check for restriction note self.assertIn("OpenRouter models restricted by", result) @patch.dict(os.environ, {"OPENROUTER_API_KEY": "test-key", "GEMINI_API_KEY": "gemini-test-key"}, clear=True) @patch("providers.registries.openrouter.OpenRouterModelRegistry") @patch.object(ModelProviderRegistry, "get_provider") def test_listmodels_shows_all_models_without_restrictions(self, mock_get_provider, mock_registry_class): """Test that listmodels shows all models when no restrictions are set.""" # Clear any cached restriction service to ensure it reads from patched environment import utils.model_restrictions utils.model_restrictions._restriction_service = None # Set up mock to return many models when no restrictions all_models = [f"provider{i // 10}/model-{i}" for i in range(50)] # Simulate 50 models from different providers self.mock_openrouter.list_models.return_value = all_models # Mock registry instance mock_registry = MagicMock() mock_registry_class.return_value = mock_registry mock_registry.resolve.return_value = None # No configs for simplicity # 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 # Create tool and execute tool = ListModelsTool() # Execute asynchronously loop = asyncio.new_event_loop() asyncio.set_event_loop(loop) result_contents = loop.run_until_complete(tool.execute({})) loop.close() # Extract text content from result result_text = result_contents[0].text # Parse JSON response import json result_json = json.loads(result_text) result = result_json["content"] # Count OpenRouter models specifically lines = result.split("\n") openrouter_section_found = False openrouter_model_count = 0 for line in lines: if "OpenRouter" in line and "✅" in line: openrouter_section_found = True elif "Custom/Local API" in line: # End of OpenRouter section break elif openrouter_section_found and line.strip().startswith("- ") and "`" in line: openrouter_model_count += 1 # After removing limits, the tool shows ALL available models (no truncation) # With 50 models from providers, we expect to see ALL of them self.assertGreaterEqual( openrouter_model_count, 30, f"Expected to see many OpenRouter models (no limits), found {openrouter_model_count}", ) # Should NOT show "and X more models available" message since we show all models now self.assertNotIn("more models available", result) # Verify list_models was called with respect_restrictions=True # (even without restrictions, we always pass True) self.mock_openrouter.list_models.assert_called_with(respect_restrictions=True) # Should NOT have restriction note when no restrictions are set self.assertNotIn("Restricted to models matching:", result) if __name__ == "__main__": unittest.main()