ff9a07a37a
feat: Azure OpenAI / Azure AI Foundry support. Models should be defined in conf/azure_models.json (or a custom path). See .env.example for environment variables or see readme. https://github.com/BeehiveInnovations/zen-mcp-server/issues/265 feat: OpenRouter / Custom Models / Azure can separately also use custom config paths now (see .env.example ) refactor: Model registry class made abstract, OpenRouter / Custom Provider / Azure OpenAI now subclass these refactor: breaking change: `is_custom` property has been removed from model_capabilities.py (and thus custom_models.json) given each models are now read from separate configuration files
369 lines
16 KiB
Python
369 lines
16 KiB
Python
"""Tests for OpenRouter provider."""
|
|
|
|
import os
|
|
from unittest.mock import Mock, patch
|
|
|
|
import pytest
|
|
|
|
from providers.openrouter import OpenRouterProvider
|
|
from providers.registry import ModelProviderRegistry
|
|
from providers.shared import ProviderType
|
|
|
|
|
|
class TestOpenRouterProvider:
|
|
"""Test cases for OpenRouter provider."""
|
|
|
|
def test_provider_initialization(self):
|
|
"""Test OpenRouter provider initialization."""
|
|
provider = OpenRouterProvider(api_key="test-key")
|
|
assert provider.api_key == "test-key"
|
|
assert provider.base_url == "https://openrouter.ai/api/v1"
|
|
assert provider.FRIENDLY_NAME == "OpenRouter"
|
|
|
|
def test_custom_headers(self):
|
|
"""Test OpenRouter custom headers."""
|
|
# Test default headers
|
|
assert "HTTP-Referer" in OpenRouterProvider.DEFAULT_HEADERS
|
|
assert "X-Title" in OpenRouterProvider.DEFAULT_HEADERS
|
|
|
|
# Test with environment variables
|
|
with patch.dict(os.environ, {"OPENROUTER_REFERER": "https://myapp.com", "OPENROUTER_TITLE": "My App"}):
|
|
from importlib import reload
|
|
|
|
import providers.openrouter
|
|
|
|
reload(providers.openrouter)
|
|
|
|
provider = providers.openrouter.OpenRouterProvider(api_key="test-key")
|
|
assert provider.DEFAULT_HEADERS["HTTP-Referer"] == "https://myapp.com"
|
|
assert provider.DEFAULT_HEADERS["X-Title"] == "My App"
|
|
|
|
def test_model_validation(self):
|
|
"""Test model validation."""
|
|
provider = OpenRouterProvider(api_key="test-key")
|
|
|
|
# OpenRouter accepts models with provider prefixes or known models
|
|
assert provider.validate_model_name("openai/gpt-4") is True
|
|
assert provider.validate_model_name("anthropic/claude-3-opus") is True
|
|
assert provider.validate_model_name("google/any-model-name") is True
|
|
assert provider.validate_model_name("groq/llama-3.1-8b") is True
|
|
|
|
# Unknown models without provider prefix are rejected
|
|
assert provider.validate_model_name("gpt-4") is False
|
|
assert provider.validate_model_name("unknown-model") is False
|
|
|
|
def test_get_capabilities(self):
|
|
"""Test capability generation."""
|
|
provider = OpenRouterProvider(api_key="test-key")
|
|
|
|
# Test with a model in the registry (using alias)
|
|
caps = provider.get_capabilities("o3")
|
|
assert caps.provider == ProviderType.OPENROUTER
|
|
assert caps.model_name == "openai/o3" # Resolved name
|
|
assert caps.friendly_name == "OpenRouter (openai/o3)"
|
|
|
|
# Test with a model not in registry - should raise error
|
|
with pytest.raises(ValueError, match="Unsupported model 'unknown-model' for provider openrouter"):
|
|
provider.get_capabilities("unknown-model")
|
|
|
|
# Test with model that has provider prefix - should get generic capabilities
|
|
caps = provider.get_capabilities("provider/unknown-model")
|
|
assert caps.provider == ProviderType.OPENROUTER
|
|
assert caps.model_name == "provider/unknown-model"
|
|
assert caps.context_window == 32_768 # Safe default
|
|
assert hasattr(caps, "_is_generic") and caps._is_generic is True
|
|
|
|
def test_model_alias_resolution(self):
|
|
"""Test model alias resolution."""
|
|
provider = OpenRouterProvider(api_key="test-key")
|
|
|
|
# Test alias resolution
|
|
assert provider._resolve_model_name("opus") == "anthropic/claude-opus-4.1"
|
|
assert provider._resolve_model_name("sonnet") == "anthropic/claude-sonnet-4.5"
|
|
assert provider._resolve_model_name("sonnet4.1") == "anthropic/claude-sonnet-4.1"
|
|
assert provider._resolve_model_name("o3") == "openai/o3"
|
|
assert provider._resolve_model_name("o3-mini") == "openai/o3-mini"
|
|
assert provider._resolve_model_name("o3mini") == "openai/o3-mini"
|
|
assert provider._resolve_model_name("o4-mini") == "openai/o4-mini"
|
|
assert provider._resolve_model_name("o4-mini") == "openai/o4-mini"
|
|
assert provider._resolve_model_name("haiku") == "anthropic/claude-3.5-haiku"
|
|
assert provider._resolve_model_name("mistral") == "mistralai/mistral-large-2411"
|
|
assert provider._resolve_model_name("deepseek") == "deepseek/deepseek-r1-0528"
|
|
assert provider._resolve_model_name("r1") == "deepseek/deepseek-r1-0528"
|
|
|
|
# Test case-insensitive
|
|
assert provider._resolve_model_name("OPUS") == "anthropic/claude-opus-4.1"
|
|
assert provider._resolve_model_name("SONNET") == "anthropic/claude-sonnet-4.5"
|
|
assert provider._resolve_model_name("O3") == "openai/o3"
|
|
assert provider._resolve_model_name("Mistral") == "mistralai/mistral-large-2411"
|
|
|
|
# Test direct model names (should pass through unchanged)
|
|
assert provider._resolve_model_name("anthropic/claude-opus-4.1") == "anthropic/claude-opus-4.1"
|
|
assert provider._resolve_model_name("openai/o3") == "openai/o3"
|
|
|
|
# Test unknown models pass through
|
|
assert provider._resolve_model_name("unknown-model") == "unknown-model"
|
|
assert provider._resolve_model_name("custom/model-v2") == "custom/model-v2"
|
|
|
|
def test_openrouter_registration(self):
|
|
"""Test OpenRouter can be registered and retrieved."""
|
|
with patch.dict(os.environ, {"OPENROUTER_API_KEY": "test-key"}):
|
|
# Clean up any existing registration
|
|
ModelProviderRegistry.unregister_provider(ProviderType.OPENROUTER)
|
|
|
|
# Register the provider
|
|
ModelProviderRegistry.register_provider(ProviderType.OPENROUTER, OpenRouterProvider)
|
|
|
|
# Retrieve and verify
|
|
provider = ModelProviderRegistry.get_provider(ProviderType.OPENROUTER)
|
|
assert provider is not None
|
|
assert isinstance(provider, OpenRouterProvider)
|
|
|
|
|
|
class TestOpenRouterAutoMode:
|
|
"""Test auto mode functionality when only OpenRouter is configured."""
|
|
|
|
def setup_method(self):
|
|
"""Store original state before each test."""
|
|
self.registry = ModelProviderRegistry()
|
|
self._original_providers = self.registry._providers.copy()
|
|
self._original_initialized = self.registry._initialized_providers.copy()
|
|
|
|
self.registry._providers.clear()
|
|
self.registry._initialized_providers.clear()
|
|
|
|
self._original_env = {}
|
|
for key in ["OPENROUTER_API_KEY", "GEMINI_API_KEY", "OPENAI_API_KEY", "DEFAULT_MODEL"]:
|
|
self._original_env[key] = os.environ.get(key)
|
|
|
|
def teardown_method(self):
|
|
"""Restore original state after each test."""
|
|
self.registry._providers.clear()
|
|
self.registry._initialized_providers.clear()
|
|
self.registry._providers.update(self._original_providers)
|
|
self.registry._initialized_providers.update(self._original_initialized)
|
|
|
|
for key, value in self._original_env.items():
|
|
if value is None:
|
|
os.environ.pop(key, None)
|
|
else:
|
|
os.environ[key] = value
|
|
|
|
@pytest.mark.no_mock_provider
|
|
def test_openrouter_only_auto_mode(self):
|
|
"""Test that auto mode works when only OpenRouter is configured."""
|
|
os.environ.pop("GEMINI_API_KEY", None)
|
|
os.environ.pop("OPENAI_API_KEY", None)
|
|
os.environ["OPENROUTER_API_KEY"] = "test-openrouter-key"
|
|
os.environ["DEFAULT_MODEL"] = "auto"
|
|
|
|
mock_registry = Mock()
|
|
model_names = [
|
|
"google/gemini-2.5-flash",
|
|
"google/gemini-2.5-pro",
|
|
"openai/o3",
|
|
"openai/o3-mini",
|
|
"anthropic/claude-opus-4.1",
|
|
"anthropic/claude-sonnet-4.1",
|
|
]
|
|
mock_registry.list_models.return_value = model_names
|
|
|
|
# Mock resolve to return a ModelCapabilities-like object for each model
|
|
def mock_resolve(model_name):
|
|
if model_name in model_names:
|
|
mock_config = Mock()
|
|
mock_config.provider = ProviderType.OPENROUTER
|
|
mock_config.aliases = [] # Empty list of aliases
|
|
mock_config.get_effective_capability_rank = Mock(return_value=50) # Add ranking method
|
|
return mock_config
|
|
return None
|
|
|
|
mock_registry.resolve.side_effect = mock_resolve
|
|
|
|
ModelProviderRegistry.register_provider(ProviderType.OPENROUTER, OpenRouterProvider)
|
|
|
|
provider = ModelProviderRegistry.get_provider(ProviderType.OPENROUTER)
|
|
assert provider is not None, "OpenRouter provider should be available with API key"
|
|
provider._registry = mock_registry
|
|
|
|
available_models = ModelProviderRegistry.get_available_models(respect_restrictions=True)
|
|
|
|
assert len(available_models) > 0, "Should find OpenRouter models in auto mode"
|
|
assert all(provider_type == ProviderType.OPENROUTER for provider_type in available_models.values())
|
|
|
|
for model in model_names:
|
|
assert model in available_models, f"Model {model} should be available"
|
|
|
|
@pytest.mark.no_mock_provider
|
|
def test_openrouter_with_restrictions(self):
|
|
"""Test that OpenRouter respects model restrictions."""
|
|
os.environ.pop("GEMINI_API_KEY", None)
|
|
os.environ.pop("OPENAI_API_KEY", None)
|
|
os.environ["OPENROUTER_API_KEY"] = "test-openrouter-key"
|
|
os.environ.pop("OPENROUTER_ALLOWED_MODELS", None)
|
|
os.environ["OPENROUTER_ALLOWED_MODELS"] = "anthropic/claude-opus-4.1,google/gemini-2.5-flash"
|
|
os.environ["DEFAULT_MODEL"] = "auto"
|
|
|
|
# Force reload to pick up new environment variable
|
|
import utils.model_restrictions
|
|
|
|
utils.model_restrictions._restriction_service = None
|
|
|
|
mock_registry = Mock()
|
|
mock_models = [
|
|
"google/gemini-2.5-flash",
|
|
"google/gemini-2.5-pro",
|
|
"anthropic/claude-opus-4.1",
|
|
"anthropic/claude-sonnet-4.1",
|
|
]
|
|
mock_registry.list_models.return_value = mock_models
|
|
|
|
# Mock the resolve method to return model configs with aliases
|
|
mock_model_config = Mock()
|
|
mock_model_config.aliases = [] # Empty aliases for simplicity
|
|
mock_model_config.get_effective_capability_rank = Mock(return_value=50) # Add ranking method
|
|
mock_registry.resolve.return_value = mock_model_config
|
|
|
|
ModelProviderRegistry.register_provider(ProviderType.OPENROUTER, OpenRouterProvider)
|
|
|
|
provider = ModelProviderRegistry.get_provider(ProviderType.OPENROUTER)
|
|
provider._registry = mock_registry
|
|
|
|
available_models = ModelProviderRegistry.get_available_models(respect_restrictions=True)
|
|
|
|
assert len(available_models) > 0, "Should have some allowed models"
|
|
|
|
expected_allowed = {"google/gemini-2.5-flash", "anthropic/claude-opus-4.1"}
|
|
|
|
assert (
|
|
set(available_models.keys()) == expected_allowed
|
|
), f"Expected {expected_allowed}, but got {set(available_models.keys())}"
|
|
|
|
@pytest.mark.no_mock_provider
|
|
def test_no_providers_fails_auto_mode(self):
|
|
"""Test that auto mode fails gracefully when no providers are available."""
|
|
os.environ.pop("GEMINI_API_KEY", None)
|
|
os.environ.pop("OPENAI_API_KEY", None)
|
|
os.environ.pop("OPENROUTER_API_KEY", None)
|
|
os.environ["DEFAULT_MODEL"] = "auto"
|
|
|
|
available_models = ModelProviderRegistry.get_available_models(respect_restrictions=True)
|
|
|
|
assert len(available_models) == 0, "Should have no models when no providers are configured"
|
|
|
|
@pytest.mark.no_mock_provider
|
|
def test_openrouter_without_registry(self):
|
|
"""Test that OpenRouter without _registry attribute doesn't crash."""
|
|
os.environ.pop("GEMINI_API_KEY", None)
|
|
os.environ.pop("OPENAI_API_KEY", None)
|
|
os.environ["OPENROUTER_API_KEY"] = "test-openrouter-key"
|
|
os.environ["DEFAULT_MODEL"] = "auto"
|
|
|
|
mock_provider_class = Mock()
|
|
mock_provider_instance = Mock(spec=["get_provider_type", "list_models", "get_all_model_capabilities"])
|
|
mock_provider_instance.get_provider_type.return_value = ProviderType.OPENROUTER
|
|
mock_provider_instance.list_models.return_value = []
|
|
mock_provider_instance.get_all_model_capabilities.return_value = {}
|
|
mock_provider_class.return_value = mock_provider_instance
|
|
|
|
ModelProviderRegistry.register_provider(ProviderType.OPENROUTER, mock_provider_class)
|
|
|
|
available_models = ModelProviderRegistry.get_available_models(respect_restrictions=True)
|
|
|
|
assert len(available_models) == 0, "Should have no models when OpenRouter has no registry"
|
|
|
|
|
|
class TestOpenRouterRegistry:
|
|
"""Test cases for OpenRouter model registry."""
|
|
|
|
def test_registry_loading(self):
|
|
"""Test registry loads models from config."""
|
|
from providers.openrouter_registry import OpenRouterModelRegistry
|
|
|
|
registry = OpenRouterModelRegistry()
|
|
|
|
# Should have loaded models
|
|
models = registry.list_models()
|
|
assert len(models) > 0
|
|
assert "anthropic/claude-opus-4.1" in models
|
|
assert "openai/o3" in models
|
|
|
|
# Should have loaded aliases
|
|
aliases = registry.list_aliases()
|
|
assert len(aliases) > 0
|
|
assert "opus" in aliases
|
|
assert "o3" in aliases
|
|
assert "sonnet" in aliases
|
|
|
|
def test_registry_capabilities(self):
|
|
"""Test registry provides correct capabilities."""
|
|
from providers.openrouter_registry import OpenRouterModelRegistry
|
|
|
|
registry = OpenRouterModelRegistry()
|
|
|
|
# Test known model
|
|
caps = registry.get_capabilities("opus")
|
|
assert caps is not None
|
|
assert caps.model_name == "anthropic/claude-opus-4.1"
|
|
assert caps.context_window == 200000 # Claude's context window
|
|
|
|
# Test using full model name
|
|
caps = registry.get_capabilities("anthropic/claude-opus-4.1")
|
|
assert caps is not None
|
|
assert caps.model_name == "anthropic/claude-opus-4.1"
|
|
|
|
# Test unknown model
|
|
caps = registry.get_capabilities("non-existent-model")
|
|
assert caps is None
|
|
|
|
def test_multiple_aliases_same_model(self):
|
|
"""Test multiple aliases pointing to same model."""
|
|
from providers.openrouter_registry import OpenRouterModelRegistry
|
|
|
|
registry = OpenRouterModelRegistry()
|
|
|
|
# All these should resolve to Claude Sonnet 4.5
|
|
sonnet_45_aliases = ["sonnet", "sonnet4.5"]
|
|
for alias in sonnet_45_aliases:
|
|
config = registry.resolve(alias)
|
|
assert config is not None
|
|
assert config.model_name == "anthropic/claude-sonnet-4.5"
|
|
|
|
# Test Sonnet 4.1 alias
|
|
config = registry.resolve("sonnet4.1")
|
|
assert config is not None
|
|
assert config.model_name == "anthropic/claude-sonnet-4.1"
|
|
|
|
|
|
class TestOpenRouterFunctionality:
|
|
"""Test OpenRouter-specific functionality."""
|
|
|
|
def test_openrouter_always_uses_correct_url(self):
|
|
"""Test that OpenRouter always uses the correct base URL."""
|
|
provider = OpenRouterProvider(api_key="test-key")
|
|
assert provider.base_url == "https://openrouter.ai/api/v1"
|
|
|
|
# Even if we try to change it, it should remain the OpenRouter URL
|
|
# (This is a characteristic of the OpenRouter provider)
|
|
provider.base_url = "http://example.com" # Try to change it
|
|
# But new instances should always use the correct URL
|
|
provider2 = OpenRouterProvider(api_key="test-key")
|
|
assert provider2.base_url == "https://openrouter.ai/api/v1"
|
|
|
|
def test_openrouter_headers_set_correctly(self):
|
|
"""Test that OpenRouter specific headers are set."""
|
|
provider = OpenRouterProvider(api_key="test-key")
|
|
|
|
# Check default headers
|
|
assert "HTTP-Referer" in provider.DEFAULT_HEADERS
|
|
assert "X-Title" in provider.DEFAULT_HEADERS
|
|
assert provider.DEFAULT_HEADERS["X-Title"] == "Zen MCP Server"
|
|
|
|
def test_openrouter_model_registry_initialized(self):
|
|
"""Test that model registry is properly initialized."""
|
|
provider = OpenRouterProvider(api_key="test-key")
|
|
|
|
# Registry should be initialized
|
|
assert hasattr(provider, "_registry")
|
|
assert provider._registry is not None
|