"""Test listmodels tool respects model restrictions."""
import asyncio
import os
import unittest
from unittest.mock import MagicMock, patch
from providers.base import ModelProvider, ProviderType
from providers.registry import ModelProviderRegistry
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
# 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"]
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.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,
"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")
# 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 "Available Models" in line and openrouter_section_found:
in_openrouter_section = True
elif in_openrouter_section and line.strip().startswith("- "):
# Extract model name from various line formats:
# - `model-name` → `full-name` (context)
# - `model-name`
line_content = line.strip()[2:] # Remove "- "
if "`" in line_content:
# Extract content between first pair of backticks
model_name = line_content.split("`")[1]
openrouter_models.append(model_name)
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 list_models was called with respect_restrictions=True
self.mock_openrouter.list_models.assert_called_with(respect_restrictions=True)
# Check for restriction note
self.assertIn("Restricted to models matching:", result)
@patch.dict(os.environ, {"OPENROUTER_API_KEY": "test-key", "GEMINI_API_KEY": "gemini-test-key"}, clear=True)
@patch("providers.openrouter_registry.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()
