my-pal-mcp-server/tests/test_openrouter_provider.py

"""Tests for OpenRouter provider."""

import os
import pytest
from unittest.mock import patch, MagicMock

from providers.base import ProviderType
from providers.openrouter import OpenRouterProvider
from providers.registry import ModelProviderRegistry


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")
        
        # Should accept any model - OpenRouter handles validation
        assert provider.validate_model_name("gpt-4") is True
        assert provider.validate_model_name("claude-3-opus") is True
        assert provider.validate_model_name("any-model-name") is True
        assert provider.validate_model_name("GPT-4") is True
        assert provider.validate_model_name("unknown-model") is True
    
    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("gpt4o")
        assert caps.provider == ProviderType.OPENROUTER
        assert caps.model_name == "openai/gpt-4o"  # Resolved name
        assert caps.friendly_name == "OpenRouter"
        
        # Test with a model not in registry - should get generic capabilities
        caps = provider.get_capabilities("unknown-model")
        assert caps.provider == ProviderType.OPENROUTER
        assert caps.model_name == "unknown-model"
        assert caps.max_tokens == 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-3-opus"
        assert provider._resolve_model_name("sonnet") == "anthropic/claude-3-sonnet"
        assert provider._resolve_model_name("gpt4o") == "openai/gpt-4o"
        assert provider._resolve_model_name("4o") == "openai/gpt-4o"
        assert provider._resolve_model_name("claude") == "anthropic/claude-3-sonnet"
        assert provider._resolve_model_name("mistral") == "mistral/mistral-large"
        assert provider._resolve_model_name("deepseek") == "deepseek/deepseek-coder"
        assert provider._resolve_model_name("coder") == "deepseek/deepseek-coder"
        
        # Test case-insensitive
        assert provider._resolve_model_name("OPUS") == "anthropic/claude-3-opus"
        assert provider._resolve_model_name("GPT4O") == "openai/gpt-4o"
        assert provider._resolve_model_name("Mistral") == "mistral/mistral-large"
        assert provider._resolve_model_name("CLAUDE") == "anthropic/claude-3-sonnet"
        
        # Test direct model names (should pass through unchanged)
        assert provider._resolve_model_name("anthropic/claude-3-opus") == "anthropic/claude-3-opus"
        assert provider._resolve_model_name("openai/gpt-4o") == "openai/gpt-4o"
        
        # 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 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-3-opus" in models
        assert "openai/gpt-4o" in models
        
        # Should have loaded aliases
        aliases = registry.list_aliases()
        assert len(aliases) > 0
        assert "opus" in aliases
        assert "gpt4o" in aliases
        assert "claude" 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-3-opus"
        assert caps.max_tokens == 200000  # Claude's context window
        
        # Test using full model name
        caps = registry.get_capabilities("anthropic/claude-3-opus")
        assert caps is not None
        assert caps.model_name == "anthropic/claude-3-opus"
        
        # 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
        sonnet_aliases = ["sonnet", "claude", "claude-sonnet", "claude3-sonnet"]
        for alias in sonnet_aliases:
            config = registry.resolve(alias)
            assert config is not None
            assert config.model_name == "anthropic/claude-3-sonnet"


class TestOpenRouterSSRFProtection:
    """Test SSRF protection for OpenRouter."""
    
    def test_url_validation_rejects_private_ips(self):
        """Test that private IPs are rejected."""
        provider = OpenRouterProvider(api_key="test-key")
        
        # List of private/dangerous IPs to test
        dangerous_urls = [
            "http://192.168.1.1/api/v1",
            "http://10.0.0.1/api/v1",
            "http://172.16.0.1/api/v1",
            "http://169.254.169.254/api/v1",  # AWS metadata
            "http://[::1]/api/v1",  # IPv6 localhost
            "http://0.0.0.0/api/v1",
        ]
        
        for url in dangerous_urls:
            with pytest.raises(ValueError, match="restricted IP|Invalid"):
                provider.base_url = url
                provider._validate_base_url()
    
    def test_url_validation_allows_public_domains(self):
        """Test that legitimate public domains are allowed."""
        provider = OpenRouterProvider(api_key="test-key")
        
        # OpenRouter's actual domain should always be allowed
        provider.base_url = "https://openrouter.ai/api/v1"
        provider._validate_base_url()  # Should not raise
    
    def test_invalid_url_schemes_rejected(self):
        """Test that non-HTTP(S) schemes are rejected."""
        provider = OpenRouterProvider(api_key="test-key")
        
        invalid_urls = [
            "ftp://example.com/api",
            "file:///etc/passwd",
            "gopher://example.com",
            "javascript:alert(1)",
        ]
        
        for url in invalid_urls:
            with pytest.raises(ValueError, match="Invalid URL scheme"):
                provider.base_url = url
                provider._validate_base_url()