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

Fixed: https://github.com/BeehiveInnovations/zen-mcp-server/issues/123

Browse Source 
Fixed restriction checks for OpenRouter
This commit is contained in:
Fahad
2025-06-23 15:23:55 +04:00
parent b4852c825f
commit e94c028a3f
9 changed files with 246 additions and 60 deletions
Download Patch File Download Diff File Expand all files Collapse all files

32
tests/test_custom_provider.py
View File

@@ -45,18 +45,32 @@ class TestCustomProvider:
def test_get_capabilities_from_registry(self):
"""Test get_capabilities returns registry capabilities when available."""
provider = CustomProvider(api_key="test-key", base_url="http://localhost:11434/v1")
# Save original environment
original_env = os.environ.get("OPENROUTER_ALLOWED_MODELS")
# Test with a model that should be in the registry (OpenRouter model) and is allowed by restrictions
capabilities = provider.get_capabilities("o3") # o3 is in OPENROUTER_ALLOWED_MODELS
try:
# Clear any restrictions
os.environ.pop("OPENROUTER_ALLOWED_MODELS", None)
assert capabilities.provider == ProviderType.OPENROUTER # o3 is an OpenRouter model (is_custom=false)
assert capabilities.context_window > 0
provider = CustomProvider(api_key="test-key", base_url="http://localhost:11434/v1")
# Test with a custom model (is_custom=true)
capabilities = provider.get_capabilities("local-llama")
assert capabilities.provider == ProviderType.CUSTOM # local-llama has is_custom=true
assert capabilities.context_window > 0
# Test with a model that should be in the registry (OpenRouter model)
capabilities = provider.get_capabilities("o3") # o3 is an OpenRouter model
assert capabilities.provider == ProviderType.OPENROUTER # o3 is an OpenRouter model (is_custom=false)
assert capabilities.context_window > 0
# Test with a custom model (is_custom=true)
capabilities = provider.get_capabilities("local-llama")
assert capabilities.provider == ProviderType.CUSTOM # local-llama has is_custom=true
assert capabilities.context_window > 0
finally:
# Restore original environment
if original_env is None:
os.environ.pop("OPENROUTER_ALLOWED_MODELS", None)
else:
os.environ["OPENROUTER_ALLOWED_MODELS"] = original_env
def test_get_capabilities_generic_fallback(self):
"""Test get_capabilities returns generic capabilities for unknown models."""

79
tests/test_parse_model_option.py Normal file
View File

@@ -0,0 +1,79 @@
"""Tests for parse_model_option function."""
from server import parse_model_option
class TestParseModelOption:
"""Test cases for model option parsing."""
def test_openrouter_free_suffix_preserved(self):
"""Test that OpenRouter :free suffix is preserved as part of model name."""
model, option = parse_model_option("openai/gpt-3.5-turbo:free")
assert model == "openai/gpt-3.5-turbo:free"
assert option is None
def test_openrouter_beta_suffix_preserved(self):
"""Test that OpenRouter :beta suffix is preserved as part of model name."""
model, option = parse_model_option("anthropic/claude-3-opus:beta")
assert model == "anthropic/claude-3-opus:beta"
assert option is None
def test_openrouter_preview_suffix_preserved(self):
"""Test that OpenRouter :preview suffix is preserved as part of model name."""
model, option = parse_model_option("google/gemini-pro:preview")
assert model == "google/gemini-pro:preview"
assert option is None
def test_ollama_tag_parsed_as_option(self):
"""Test that Ollama tags are parsed as options."""
model, option = parse_model_option("llama3.2:latest")
assert model == "llama3.2"
assert option == "latest"
def test_consensus_stance_parsed_as_option(self):
"""Test that consensus stances are parsed as options."""
model, option = parse_model_option("o3:for")
assert model == "o3"
assert option == "for"
model, option = parse_model_option("gemini-2.5-pro:against")
assert model == "gemini-2.5-pro"
assert option == "against"
def test_openrouter_unknown_suffix_parsed_as_option(self):
"""Test that unknown suffixes on OpenRouter models are parsed as options."""
model, option = parse_model_option("openai/gpt-4:custom-tag")
assert model == "openai/gpt-4"
assert option == "custom-tag"
def test_plain_model_name(self):
"""Test plain model names without colons."""
model, option = parse_model_option("gpt-4")
assert model == "gpt-4"
assert option is None
def test_url_not_parsed(self):
"""Test that URLs are not parsed for options."""
model, option = parse_model_option("http://localhost:8080")
assert model == "http://localhost:8080"
assert option is None
def test_whitespace_handling(self):
"""Test that whitespace is properly stripped."""
model, option = parse_model_option(" openai/gpt-3.5-turbo:free ")
assert model == "openai/gpt-3.5-turbo:free"
assert option is None
model, option = parse_model_option(" llama3.2 : latest ")
assert model == "llama3.2"
assert option == "latest"
def test_case_insensitive_suffix_matching(self):
"""Test that OpenRouter suffix matching is case-insensitive."""
model, option = parse_model_option("openai/gpt-3.5-turbo:FREE")
assert model == "openai/gpt-3.5-turbo:FREE" # Original case preserved
assert option is None
model, option = parse_model_option("openai/gpt-3.5-turbo:Free")
assert model == "openai/gpt-3.5-turbo:Free" # Original case preserved
assert option is None

34
tests/test_provider_routing_bugs.py
View File

@@ -58,7 +58,13 @@ class TestProviderRoutingBugs:
"""
# Save original environment
original_env = {}
for key in ["GEMINI_API_KEY", "OPENAI_API_KEY", "XAI_API_KEY", "OPENROUTER_API_KEY"]:
for key in [
"GEMINI_API_KEY",
"OPENAI_API_KEY",
"XAI_API_KEY",
"OPENROUTER_API_KEY",
"OPENROUTER_ALLOWED_MODELS",
]:
original_env[key] = os.environ.get(key)
try:
@@ -66,6 +72,7 @@ class TestProviderRoutingBugs:
os.environ.pop("GEMINI_API_KEY", None) # No Google API key
os.environ.pop("OPENAI_API_KEY", None)
os.environ.pop("XAI_API_KEY", None)
os.environ.pop("OPENROUTER_ALLOWED_MODELS", None) # Clear any restrictions
os.environ["OPENROUTER_API_KEY"] = "test-openrouter-key"
# Register only OpenRouter provider (like in server.py:configure_providers)
@@ -113,12 +120,24 @@ class TestProviderRoutingBugs:
"""
# Save original environment
original_env = {}
for key in ["GEMINI_API_KEY", "OPENAI_API_KEY", "XAI_API_KEY", "OPENROUTER_API_KEY"]:
for key in [
"GEMINI_API_KEY",
"OPENAI_API_KEY",
"XAI_API_KEY",
"OPENROUTER_API_KEY",
"OPENROUTER_ALLOWED_MODELS",
]:
original_env[key] = os.environ.get(key)
try:
# Set up scenario: NO API keys at all
for key in ["GEMINI_API_KEY", "OPENAI_API_KEY", "XAI_API_KEY", "OPENROUTER_API_KEY"]:
for key in [
"GEMINI_API_KEY",
"OPENAI_API_KEY",
"XAI_API_KEY",
"OPENROUTER_API_KEY",
"OPENROUTER_ALLOWED_MODELS",
]:
os.environ.pop(key, None)
# Create tool to test fallback logic
@@ -151,7 +170,13 @@ class TestProviderRoutingBugs:
"""
# Save original environment
original_env = {}
for key in ["GEMINI_API_KEY", "OPENAI_API_KEY", "XAI_API_KEY", "OPENROUTER_API_KEY"]:
for key in [
"GEMINI_API_KEY",
"OPENAI_API_KEY",
"XAI_API_KEY",
"OPENROUTER_API_KEY",
"OPENROUTER_ALLOWED_MODELS",
]:
original_env[key] = os.environ.get(key)
try:
@@ -160,6 +185,7 @@ class TestProviderRoutingBugs:
os.environ["OPENAI_API_KEY"] = "test-openai-key"
os.environ["OPENROUTER_API_KEY"] = "test-openrouter-key"
os.environ.pop("XAI_API_KEY", None)
os.environ.pop("OPENROUTER_ALLOWED_MODELS", None) # Clear any restrictions
# Register providers in priority order (like server.py)
from providers.gemini import GeminiModelProvider

104
tests/test_workflow_metadata.py
View File

@@ -48,7 +48,13 @@ class TestWorkflowMetadata:
"""
# Save original environment
original_env = {}
for key in ["GEMINI_API_KEY", "OPENAI_API_KEY", "XAI_API_KEY", "OPENROUTER_API_KEY"]:
for key in [
"GEMINI_API_KEY",
"OPENAI_API_KEY",
"XAI_API_KEY",
"OPENROUTER_API_KEY",
"OPENROUTER_ALLOWED_MODELS",
]:
original_env[key] = os.environ.get(key)
try:
@@ -56,6 +62,7 @@ class TestWorkflowMetadata:
os.environ.pop("GEMINI_API_KEY", None)
os.environ.pop("OPENAI_API_KEY", None)
os.environ.pop("XAI_API_KEY", None)
os.environ.pop("OPENROUTER_ALLOWED_MODELS", None) # Clear any restrictions
os.environ["OPENROUTER_API_KEY"] = "test-openrouter-key"
# Register OpenRouter provider
@@ -124,7 +131,13 @@ class TestWorkflowMetadata:
"""
# Save original environment
original_env = {}
for key in ["GEMINI_API_KEY", "OPENAI_API_KEY", "XAI_API_KEY", "OPENROUTER_API_KEY"]:
for key in [
"GEMINI_API_KEY",
"OPENAI_API_KEY",
"XAI_API_KEY",
"OPENROUTER_API_KEY",
"OPENROUTER_ALLOWED_MODELS",
]:
original_env[key] = os.environ.get(key)
try:
@@ -132,6 +145,7 @@ class TestWorkflowMetadata:
os.environ.pop("GEMINI_API_KEY", None)
os.environ.pop("OPENAI_API_KEY", None)
os.environ.pop("XAI_API_KEY", None)
os.environ.pop("OPENROUTER_ALLOWED_MODELS", None) # Clear any restrictions
os.environ["OPENROUTER_API_KEY"] = "test-openrouter-key"
# Register OpenRouter provider
@@ -182,43 +196,60 @@ class TestWorkflowMetadata:
"""
Test that workflow tools handle metadata gracefully when model context is missing.
"""
# Create debug tool
debug_tool = DebugIssueTool()
# Save original environment
original_env = {}
for key in ["OPENROUTER_ALLOWED_MODELS"]:
original_env[key] = os.environ.get(key)
# Create arguments without model context (fallback scenario)
arguments = {
"step": "Test step without model context",
"step_number": 1,
"total_steps": 1,
"next_step_required": False,
"findings": "Test findings",
"model": "flash",
"confidence": "low",
# No _model_context or _resolved_model_name
}
try:
# Clear any restrictions
os.environ.pop("OPENROUTER_ALLOWED_MODELS", None)
# Execute the workflow tool
import asyncio
# Create debug tool
debug_tool = DebugIssueTool()
result = asyncio.run(debug_tool.execute_workflow(arguments))
# Create arguments without model context (fallback scenario)
arguments = {
"step": "Test step without model context",
"step_number": 1,
"total_steps": 1,
"next_step_required": False,
"findings": "Test findings",
"model": "flash",
"confidence": "low",
# No _model_context or _resolved_model_name
}
# Parse the JSON response
assert len(result) == 1
response_text = result[0].text
response_data = json.loads(response_text)
# Execute the workflow tool
import asyncio
# Verify metadata is still present with fallback values
assert "metadata" in response_data, "Workflow response should include metadata even in fallback"
metadata = response_data["metadata"]
result = asyncio.run(debug_tool.execute_workflow(arguments))
# Verify fallback metadata
assert "tool_name" in metadata, "Fallback metadata should include tool_name"
assert "model_used" in metadata, "Fallback metadata should include model_used"
assert "provider_used" in metadata, "Fallback metadata should include provider_used"
# Parse the JSON response
assert len(result) == 1
response_text = result[0].text
response_data = json.loads(response_text)
assert metadata["tool_name"] == "debug", "tool_name should be 'debug'"
assert metadata["model_used"] == "flash", "model_used should be from request"
assert metadata["provider_used"] == "unknown", "provider_used should be 'unknown' in fallback"
# Verify metadata is still present with fallback values
assert "metadata" in response_data, "Workflow response should include metadata even in fallback"
metadata = response_data["metadata"]
# Verify fallback metadata
assert "tool_name" in metadata, "Fallback metadata should include tool_name"
assert "model_used" in metadata, "Fallback metadata should include model_used"
assert "provider_used" in metadata, "Fallback metadata should include provider_used"
assert metadata["tool_name"] == "debug", "tool_name should be 'debug'"
assert metadata["model_used"] == "flash", "model_used should be from request"
assert metadata["provider_used"] == "unknown", "provider_used should be 'unknown' in fallback"
finally:
# Restore original environment
for key, value in original_env.items():
if value is None:
os.environ.pop(key, None)
else:
os.environ[key] = value
@pytest.mark.no_mock_provider
def test_workflow_metadata_preserves_existing_response_fields(self):
@@ -227,7 +258,13 @@ class TestWorkflowMetadata:
"""
# Save original environment
original_env = {}
for key in ["GEMINI_API_KEY", "OPENAI_API_KEY", "XAI_API_KEY", "OPENROUTER_API_KEY"]:
for key in [
"GEMINI_API_KEY",
"OPENAI_API_KEY",
"XAI_API_KEY",
"OPENROUTER_API_KEY",
"OPENROUTER_ALLOWED_MODELS",
]:
original_env[key] = os.environ.get(key)
try:
@@ -235,6 +272,7 @@ class TestWorkflowMetadata:
os.environ.pop("GEMINI_API_KEY", None)
os.environ.pop("OPENAI_API_KEY", None)
os.environ.pop("XAI_API_KEY", None)
os.environ.pop("OPENROUTER_ALLOWED_MODELS", None) # Clear any restrictions
os.environ["OPENROUTER_API_KEY"] = "test-openrouter-key"
# Register OpenRouter provider