torbjorn/my-pal-mcp-server
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Apply black formatting to test_openrouter_provider.py

Browse Source
This commit is contained in:
Nikolai Ugelvik
2025-06-14 19:33:20 +02:00
parent be2612752a
commit 0eeea3dd67
1 changed files with 59 additions and 60 deletions
Download Patch File Download Diff File Expand all files Collapse all files

71
tests/test_openrouter_provider.py
View File

@@ -125,7 +125,7 @@ class TestOpenRouterAutoMode:
self.registry._initialized_providers.clear() self.registry._initialized_providers.clear()
self._original_env = {} self._original_env = {}
for key in ['OPENROUTER_API_KEY', 'GEMINI_API_KEY', 'OPENAI_API_KEY', 'DEFAULT_MODEL']: for key in ["OPENROUTER_API_KEY", "GEMINI_API_KEY", "OPENAI_API_KEY", "DEFAULT_MODEL"]:
self._original_env[key] = os.environ.get(key) self._original_env[key] = os.environ.get(key)
def teardown_method(self): def teardown_method(self):
@@ -144,19 +144,19 @@ class TestOpenRouterAutoMode:
@pytest.mark.no_mock_provider @pytest.mark.no_mock_provider
def test_openrouter_only_auto_mode(self): def test_openrouter_only_auto_mode(self):
"""Test that auto mode works when only OpenRouter is configured.""" """Test that auto mode works when only OpenRouter is configured."""
os.environ.pop('GEMINI_API_KEY', None) os.environ.pop("GEMINI_API_KEY", None)
os.environ.pop('OPENAI_API_KEY', None) os.environ.pop("OPENAI_API_KEY", None)
os.environ['OPENROUTER_API_KEY'] = 'test-openrouter-key' os.environ["OPENROUTER_API_KEY"] = "test-openrouter-key"
os.environ['DEFAULT_MODEL'] = 'auto' os.environ["DEFAULT_MODEL"] = "auto"
mock_registry = Mock() mock_registry = Mock()
mock_registry.list_models.return_value = [ mock_registry.list_models.return_value = [
'google/gemini-2.5-flash-preview-05-20', "google/gemini-2.5-flash-preview-05-20",
'google/gemini-2.5-pro-preview-06-05', "google/gemini-2.5-pro-preview-06-05",
'openai/o3', "openai/o3",
'openai/o3-mini', "openai/o3-mini",
'anthropic/claude-3-opus', "anthropic/claude-3-opus",
'anthropic/claude-3-sonnet' "anthropic/claude-3-sonnet",
] ]
ModelProviderRegistry.register_provider(ProviderType.OPENROUTER, OpenRouterProvider) ModelProviderRegistry.register_provider(ProviderType.OPENROUTER, OpenRouterProvider)
@@ -177,23 +177,24 @@ class TestOpenRouterAutoMode:
@pytest.mark.no_mock_provider @pytest.mark.no_mock_provider
def test_openrouter_with_restrictions(self): def test_openrouter_with_restrictions(self):
"""Test that OpenRouter respects model restrictions.""" """Test that OpenRouter respects model restrictions."""
os.environ.pop('GEMINI_API_KEY', None) os.environ.pop("GEMINI_API_KEY", None)
os.environ.pop('OPENAI_API_KEY', None) os.environ.pop("OPENAI_API_KEY", None)
os.environ['OPENROUTER_API_KEY'] = 'test-openrouter-key' os.environ["OPENROUTER_API_KEY"] = "test-openrouter-key"
os.environ.pop('OPENROUTER_ALLOWED_MODELS', None) os.environ.pop("OPENROUTER_ALLOWED_MODELS", None)
os.environ['OPENROUTER_ALLOWED_MODELS'] = 'anthropic/claude-3-opus,google/gemini-2.5-flash-preview-05-20' os.environ["OPENROUTER_ALLOWED_MODELS"] = "anthropic/claude-3-opus,google/gemini-2.5-flash-preview-05-20"
os.environ['DEFAULT_MODEL'] = 'auto' os.environ["DEFAULT_MODEL"] = "auto"
# Force reload to pick up new environment variable # Force reload to pick up new environment variable
import utils.model_restrictions import utils.model_restrictions
utils.model_restrictions._restriction_service = None utils.model_restrictions._restriction_service = None
mock_registry = Mock() mock_registry = Mock()
mock_registry.list_models.return_value = [ mock_registry.list_models.return_value = [
'google/gemini-2.5-flash-preview-05-20', "google/gemini-2.5-flash-preview-05-20",
'google/gemini-2.5-pro-preview-06-05', "google/gemini-2.5-pro-preview-06-05",
'anthropic/claude-3-opus', "anthropic/claude-3-opus",
'anthropic/claude-3-sonnet' "anthropic/claude-3-sonnet",
] ]
ModelProviderRegistry.register_provider(ProviderType.OPENROUTER, OpenRouterProvider) ModelProviderRegistry.register_provider(ProviderType.OPENROUTER, OpenRouterProvider)
@@ -205,21 +206,19 @@ class TestOpenRouterAutoMode:
assert len(available_models) > 0, "Should have some allowed models" assert len(available_models) > 0, "Should have some allowed models"
expected_allowed = { expected_allowed = {"google/gemini-2.5-flash-preview-05-20", "anthropic/claude-3-opus"}
'google/gemini-2.5-flash-preview-05-20',
'anthropic/claude-3-opus'
}
assert set(available_models.keys()) == expected_allowed, \ assert (
f"Expected {expected_allowed}, but got {set(available_models.keys())}" set(available_models.keys()) == expected_allowed
), f"Expected {expected_allowed}, but got {set(available_models.keys())}"
@pytest.mark.no_mock_provider @pytest.mark.no_mock_provider
def test_no_providers_fails_auto_mode(self): def test_no_providers_fails_auto_mode(self):
"""Test that auto mode fails gracefully when no providers are available.""" """Test that auto mode fails gracefully when no providers are available."""
os.environ.pop('GEMINI_API_KEY', None) os.environ.pop("GEMINI_API_KEY", None)
os.environ.pop('OPENAI_API_KEY', None) os.environ.pop("OPENAI_API_KEY", None)
os.environ.pop('OPENROUTER_API_KEY', None) os.environ.pop("OPENROUTER_API_KEY", None)
os.environ['DEFAULT_MODEL'] = 'auto' os.environ["DEFAULT_MODEL"] = "auto"
available_models = ModelProviderRegistry.get_available_models(respect_restrictions=True) available_models = ModelProviderRegistry.get_available_models(respect_restrictions=True)
@@ -228,13 +227,13 @@ class TestOpenRouterAutoMode:
@pytest.mark.no_mock_provider @pytest.mark.no_mock_provider
def test_openrouter_without_registry(self): def test_openrouter_without_registry(self):
"""Test that OpenRouter without _registry attribute doesn't crash.""" """Test that OpenRouter without _registry attribute doesn't crash."""
os.environ.pop('GEMINI_API_KEY', None) os.environ.pop("GEMINI_API_KEY", None)
os.environ.pop('OPENAI_API_KEY', None) os.environ.pop("OPENAI_API_KEY", None)
os.environ['OPENROUTER_API_KEY'] = 'test-openrouter-key' os.environ["OPENROUTER_API_KEY"] = "test-openrouter-key"
os.environ['DEFAULT_MODEL'] = 'auto' os.environ["DEFAULT_MODEL"] = "auto"
mock_provider_class = Mock() mock_provider_class = Mock()
mock_provider_instance = Mock(spec=['get_provider_type']) mock_provider_instance = Mock(spec=["get_provider_type"])
mock_provider_instance.get_provider_type.return_value = ProviderType.OPENROUTER mock_provider_instance.get_provider_type.return_value = ProviderType.OPENROUTER
mock_provider_class.return_value = mock_provider_instance mock_provider_class.return_value = mock_provider_instance