forked from cardosofelipe/fast-next-template
Felipe Cardoso
6e8b0b022a
Implements complete LLM Gateway MCP Server with: - FastMCP server with 4 tools: chat_completion, list_models, get_usage, count_tokens - LiteLLM Router with multi-provider failover chains - Circuit breaker pattern for fault tolerance - Redis-based cost tracking per project/agent - Comprehensive test suite (209 tests, 92% coverage) Model groups defined per ADR-004: - reasoning: claude-opus-4 → gpt-4.1 → gemini-2.5-pro - code: claude-sonnet-4 → gpt-4.1 → deepseek-coder - fast: claude-haiku → gpt-4.1-mini → gemini-2.0-flash 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
309 lines
10 KiB
Python
309 lines
10 KiB
Python
"""
|
|
Tests for providers module.
|
|
"""
|
|
|
|
import os
|
|
from unittest.mock import patch
|
|
|
|
import pytest
|
|
|
|
from config import Settings
|
|
from models import MODEL_CONFIGS, ModelGroup, Provider
|
|
from providers import (
|
|
LLMProvider,
|
|
build_fallback_config,
|
|
build_model_list,
|
|
configure_litellm,
|
|
get_available_model_groups,
|
|
get_available_models,
|
|
get_provider,
|
|
reset_provider,
|
|
)
|
|
|
|
|
|
@pytest.fixture
|
|
def full_settings() -> Settings:
|
|
"""Settings with all providers configured."""
|
|
return Settings(
|
|
anthropic_api_key="test-anthropic-key",
|
|
openai_api_key="test-openai-key",
|
|
google_api_key="test-google-key",
|
|
alibaba_api_key="test-alibaba-key",
|
|
deepseek_api_key="test-deepseek-key",
|
|
litellm_timeout=60,
|
|
litellm_cache_enabled=False,
|
|
)
|
|
|
|
|
|
@pytest.fixture
|
|
def partial_settings() -> Settings:
|
|
"""Settings with only some providers configured."""
|
|
return Settings(
|
|
anthropic_api_key="test-anthropic-key",
|
|
openai_api_key=None,
|
|
google_api_key=None,
|
|
alibaba_api_key=None,
|
|
deepseek_api_key=None,
|
|
)
|
|
|
|
|
|
@pytest.fixture
|
|
def empty_settings() -> Settings:
|
|
"""Settings with no providers configured."""
|
|
return Settings(
|
|
anthropic_api_key=None,
|
|
openai_api_key=None,
|
|
google_api_key=None,
|
|
alibaba_api_key=None,
|
|
deepseek_api_key=None,
|
|
)
|
|
|
|
|
|
class TestConfigureLiteLLM:
|
|
"""Tests for configure_litellm function."""
|
|
|
|
def test_sets_api_keys(self, full_settings: Settings) -> None:
|
|
"""Test that API keys are set in environment."""
|
|
with patch.dict(os.environ, {}, clear=True):
|
|
configure_litellm(full_settings)
|
|
|
|
assert os.environ.get("ANTHROPIC_API_KEY") == "test-anthropic-key"
|
|
assert os.environ.get("OPENAI_API_KEY") == "test-openai-key"
|
|
assert os.environ.get("GEMINI_API_KEY") == "test-google-key"
|
|
|
|
def test_skips_none_keys(self, partial_settings: Settings) -> None:
|
|
"""Test that None keys are not set."""
|
|
with patch.dict(os.environ, {}, clear=True):
|
|
configure_litellm(partial_settings)
|
|
|
|
assert os.environ.get("ANTHROPIC_API_KEY") == "test-anthropic-key"
|
|
assert "OPENAI_API_KEY" not in os.environ
|
|
|
|
|
|
class TestBuildModelList:
|
|
"""Tests for build_model_list function."""
|
|
|
|
def test_build_with_all_providers(self, full_settings: Settings) -> None:
|
|
"""Test building model list with all providers."""
|
|
model_list = build_model_list(full_settings)
|
|
|
|
assert len(model_list) > 0
|
|
|
|
# Check structure
|
|
for entry in model_list:
|
|
assert "model_name" in entry
|
|
assert "litellm_params" in entry
|
|
assert "model" in entry["litellm_params"]
|
|
assert "timeout" in entry["litellm_params"]
|
|
|
|
def test_build_with_partial_providers(self, partial_settings: Settings) -> None:
|
|
"""Test building model list with partial providers."""
|
|
model_list = build_model_list(partial_settings)
|
|
|
|
# Should only include Anthropic models
|
|
providers = set()
|
|
for entry in model_list:
|
|
model_name = entry["model_name"]
|
|
config = MODEL_CONFIGS.get(model_name)
|
|
if config:
|
|
providers.add(config.provider)
|
|
|
|
assert Provider.ANTHROPIC in providers
|
|
assert Provider.OPENAI not in providers
|
|
|
|
def test_build_with_no_providers(self, empty_settings: Settings) -> None:
|
|
"""Test building model list with no providers."""
|
|
model_list = build_model_list(empty_settings)
|
|
|
|
assert len(model_list) == 0
|
|
|
|
def test_build_includes_timeout(self, full_settings: Settings) -> None:
|
|
"""Test that model entries include timeout."""
|
|
model_list = build_model_list(full_settings)
|
|
|
|
for entry in model_list:
|
|
assert entry["litellm_params"]["timeout"] == 60
|
|
|
|
|
|
class TestBuildFallbackConfig:
|
|
"""Tests for build_fallback_config function."""
|
|
|
|
def test_build_fallbacks_full(self, full_settings: Settings) -> None:
|
|
"""Test building fallback config with all providers."""
|
|
fallbacks = build_fallback_config(full_settings)
|
|
|
|
assert len(fallbacks) > 0
|
|
|
|
# Primary models should have fallbacks
|
|
for _primary, chain in fallbacks.items():
|
|
assert isinstance(chain, list)
|
|
assert len(chain) > 0
|
|
|
|
def test_build_fallbacks_partial(self, partial_settings: Settings) -> None:
|
|
"""Test building fallback config with partial providers."""
|
|
fallbacks = build_fallback_config(partial_settings)
|
|
|
|
# With only Anthropic, there should be no fallbacks
|
|
# (fallbacks require at least 2 available models)
|
|
for primary, chain in fallbacks.items():
|
|
# All models in chain should be from Anthropic
|
|
for model in [primary] + chain:
|
|
config = MODEL_CONFIGS.get(model)
|
|
if config:
|
|
assert config.provider == Provider.ANTHROPIC
|
|
|
|
|
|
class TestGetAvailableModels:
|
|
"""Tests for get_available_models function."""
|
|
|
|
def test_get_available_full(self, full_settings: Settings) -> None:
|
|
"""Test getting available models with all providers."""
|
|
models = get_available_models(full_settings)
|
|
|
|
assert len(models) > 0
|
|
assert "claude-opus-4" in models
|
|
assert "gpt-4.1" in models
|
|
|
|
def test_get_available_partial(self, partial_settings: Settings) -> None:
|
|
"""Test getting available models with partial providers."""
|
|
models = get_available_models(partial_settings)
|
|
|
|
assert "claude-opus-4" in models
|
|
assert "gpt-4.1" not in models
|
|
|
|
def test_get_available_empty(self, empty_settings: Settings) -> None:
|
|
"""Test getting available models with no providers."""
|
|
models = get_available_models(empty_settings)
|
|
|
|
assert len(models) == 0
|
|
|
|
|
|
class TestGetAvailableModelGroups:
|
|
"""Tests for get_available_model_groups function."""
|
|
|
|
def test_get_groups_full(self, full_settings: Settings) -> None:
|
|
"""Test getting groups with all providers."""
|
|
groups = get_available_model_groups(full_settings)
|
|
|
|
assert len(groups) == len(ModelGroup)
|
|
assert ModelGroup.REASONING in groups
|
|
assert len(groups[ModelGroup.REASONING]) > 0
|
|
|
|
def test_get_groups_partial(self, partial_settings: Settings) -> None:
|
|
"""Test getting groups with partial providers."""
|
|
groups = get_available_model_groups(partial_settings)
|
|
|
|
# Only Anthropic models should be available
|
|
for _group, models in groups.items():
|
|
for model in models:
|
|
config = MODEL_CONFIGS.get(model)
|
|
if config:
|
|
assert config.provider == Provider.ANTHROPIC
|
|
|
|
|
|
class TestLLMProvider:
|
|
"""Tests for LLMProvider class."""
|
|
|
|
def test_initialization(self, full_settings: Settings) -> None:
|
|
"""Test provider initialization."""
|
|
provider = LLMProvider(settings=full_settings)
|
|
|
|
assert provider._initialized is False
|
|
assert provider._router is None
|
|
|
|
def test_initialize(self, full_settings: Settings) -> None:
|
|
"""Test provider initialize."""
|
|
with patch("providers.Router") as mock_router:
|
|
provider = LLMProvider(settings=full_settings)
|
|
provider.initialize()
|
|
|
|
assert provider._initialized is True
|
|
mock_router.assert_called_once()
|
|
|
|
def test_initialize_idempotent(self, full_settings: Settings) -> None:
|
|
"""Test that initialize is idempotent."""
|
|
with patch("providers.Router") as mock_router:
|
|
provider = LLMProvider(settings=full_settings)
|
|
provider.initialize()
|
|
provider.initialize()
|
|
|
|
# Should only be called once
|
|
assert mock_router.call_count == 1
|
|
|
|
def test_initialize_no_providers(self, empty_settings: Settings) -> None:
|
|
"""Test initialization with no providers."""
|
|
provider = LLMProvider(settings=empty_settings)
|
|
provider.initialize()
|
|
|
|
assert provider._initialized is True
|
|
assert provider._router is None
|
|
|
|
def test_router_property(self, full_settings: Settings) -> None:
|
|
"""Test router property triggers initialization."""
|
|
with patch("providers.Router"):
|
|
provider = LLMProvider(settings=full_settings)
|
|
_ = provider.router
|
|
|
|
assert provider._initialized is True
|
|
|
|
def test_is_available(self, full_settings: Settings) -> None:
|
|
"""Test is_available property."""
|
|
with patch("providers.Router"):
|
|
provider = LLMProvider(settings=full_settings)
|
|
assert provider.is_available is True
|
|
|
|
def test_is_not_available(self, empty_settings: Settings) -> None:
|
|
"""Test is_available when no providers."""
|
|
provider = LLMProvider(settings=empty_settings)
|
|
assert provider.is_available is False
|
|
|
|
def test_get_model_config(self, full_settings: Settings) -> None:
|
|
"""Test getting model config."""
|
|
provider = LLMProvider(settings=full_settings)
|
|
|
|
config = provider.get_model_config("claude-opus-4")
|
|
assert config is not None
|
|
assert config.name == "claude-opus-4"
|
|
|
|
assert provider.get_model_config("nonexistent") is None
|
|
|
|
def test_get_available_models(self, full_settings: Settings) -> None:
|
|
"""Test getting available models."""
|
|
provider = LLMProvider(settings=full_settings)
|
|
models = provider.get_available_models()
|
|
|
|
assert "claude-opus-4" in models
|
|
assert "gpt-4.1" in models
|
|
|
|
def test_is_model_available(self, full_settings: Settings) -> None:
|
|
"""Test checking model availability."""
|
|
provider = LLMProvider(settings=full_settings)
|
|
|
|
assert provider.is_model_available("claude-opus-4") is True
|
|
assert provider.is_model_available("nonexistent") is False
|
|
|
|
|
|
class TestGlobalProvider:
|
|
"""Tests for global provider functions."""
|
|
|
|
def test_get_provider(self) -> None:
|
|
"""Test getting global provider."""
|
|
reset_provider()
|
|
provider = get_provider()
|
|
assert isinstance(provider, LLMProvider)
|
|
|
|
def test_get_provider_singleton(self) -> None:
|
|
"""Test provider is singleton."""
|
|
reset_provider()
|
|
provider1 = get_provider()
|
|
provider2 = get_provider()
|
|
assert provider1 is provider2
|
|
|
|
def test_reset_provider(self) -> None:
|
|
"""Test resetting global provider."""
|
|
reset_provider()
|
|
provider1 = get_provider()
|
|
reset_provider()
|
|
provider2 = get_provider()
|
|
assert provider1 is not provider2
|