71 lines
No EOL
2.6 KiB
Python
71 lines
No EOL
2.6 KiB
Python
"""Test configuration loading from environment."""
|
|
|
|
import pytest
|
|
from router.config import RouterConfig
|
|
|
|
|
|
class TestRouterConfig:
|
|
"""Test RouterConfig functionality."""
|
|
|
|
def test_config_from_env(self, api_keys):
|
|
"""Test that config loads API keys from environment."""
|
|
config = RouterConfig.from_env()
|
|
|
|
# Check API keys are loaded
|
|
assert config.cohere_api_key is not None, "Cohere API key not loaded"
|
|
assert config.gemini_api_key is not None, "Gemini API key not loaded"
|
|
assert config.openai_api_key is not None, "OpenAI API key not loaded"
|
|
|
|
# Check Ollama configuration
|
|
assert config.ollama_base_url == api_keys["ollama_base_url"]
|
|
|
|
def test_get_api_key(self):
|
|
"""Test getting API keys by provider."""
|
|
config = RouterConfig.from_env()
|
|
|
|
assert config.get_api_key("cohere") is not None
|
|
assert config.get_api_key("gemini") is not None
|
|
assert config.get_api_key("openai") is not None
|
|
assert config.get_api_key("unknown") is None
|
|
|
|
def test_cost_calculation(self):
|
|
"""Test cost calculation methods."""
|
|
config = RouterConfig.from_env()
|
|
config.track_costs = True
|
|
|
|
# Test token cost calculation
|
|
cost = config.calculate_cost("gpt-4o", input_tokens=1000, output_tokens=500)
|
|
assert cost is not None
|
|
assert cost > 0
|
|
|
|
# Test embedding cost
|
|
embed_cost = config.calculate_embed_cost("text-embedding-004", num_tokens=1000)
|
|
assert embed_cost is not None
|
|
assert embed_cost >= 0
|
|
|
|
# Test rerank cost
|
|
rerank_cost = config.calculate_rerank_cost("rerank-english-v3.0", num_searches=10)
|
|
assert rerank_cost is not None
|
|
assert rerank_cost > 0
|
|
|
|
def test_ollama_models_config(self):
|
|
"""Test Ollama models configuration."""
|
|
config = RouterConfig.from_env()
|
|
|
|
assert "mxbai-embed-large:latest" in config.ollama_embedding_models
|
|
assert "nomic-embed-text:latest" in config.ollama_embedding_models
|
|
assert len(config.ollama_embedding_models) >= 2
|
|
|
|
def test_config_to_dict(self):
|
|
"""Test config serialization."""
|
|
config = RouterConfig.from_env()
|
|
config_dict = config.to_dict()
|
|
|
|
# Check that API keys are masked
|
|
assert config_dict["openai_api_key"] == "***"
|
|
assert config_dict["cohere_api_key"] == "***"
|
|
assert config_dict["gemini_api_key"] == "***"
|
|
|
|
# Check other settings are present
|
|
assert "default_temperature" in config_dict
|
|
assert "track_costs" in config_dict |