Files
zzrouter/backend/tests/test_key_selector.py
tech d080e67c29 test: add comprehensive test suite with 89% coverage
- Add tests for all backend modules: config, database, models, services,
  routes, middleware, and providers
- Separate dev dependencies using dependency-groups
- Update CLAUDE.md with test commands and project structure
- Add .coverage and .pytest_cache to gitignore

88 tests covering:
- Authentication and authorization
- Model routing and key selection
- Async logging with batch processing
- All API endpoints (chat, openai, anthropic, health)
- LiteLLM wrapper (mocked external calls)

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
2026-03-11 19:12:57 +08:00

93 lines
3.4 KiB
Python

"""Tests for key selector service."""
import pytest
from unittest.mock import patch
from app.services.key_selector import KeySelector
from app.models import ProviderKey
class TestKeySelector:
"""Test KeySelector."""
@pytest.mark.asyncio
async def test_get_next_key_single_key(self, mock_db):
"""Test getting next key with single key configured."""
selector = KeySelector()
with patch("app.services.key_selector.db_connection", mock_db):
# First call
key1 = await selector.get_next_key(1)
assert key1 is not None
assert isinstance(key1, ProviderKey)
@pytest.mark.asyncio
async def test_get_next_key_round_robin(self, mock_db):
"""Test round-robin key selection."""
selector = KeySelector()
selector.reset_index(1) # Reset to start fresh
with patch("app.services.key_selector.db_connection", mock_db):
keys = []
for _ in range(4): # More calls than keys (2 keys)
key = await selector.get_next_key(1)
if key:
keys.append(key.key)
# Should cycle through keys
assert len(keys) == 4
# First and third should be same, second and fourth should be same
assert keys[0] == keys[2]
assert keys[1] == keys[3]
assert keys[0] != keys[1]
@pytest.mark.asyncio
async def test_get_next_key_no_keys_configured(self, mock_db):
"""Test when no keys are configured for provider."""
selector = KeySelector()
with patch("app.services.key_selector.db_connection", mock_db):
key = await selector.get_next_key(999) # Non-existent provider
assert key is None
@pytest.mark.asyncio
async def test_reset_index(self, mock_db):
"""Test resetting round-robin index."""
selector = KeySelector()
selector._indices[1] = 5
selector.reset_index(1)
assert 1 not in selector._indices
@pytest.mark.asyncio
async def test_reset_index_nonexistent(self):
"""Test resetting index for non-existent provider."""
selector = KeySelector()
selector.reset_index(999) # Should not raise
assert 999 not in selector._indices
@pytest.mark.asyncio
async def test_concurrent_access(self, mock_db):
"""Test thread-safe concurrent access."""
import asyncio
selector = KeySelector()
selector.reset_index(1)
with patch("app.services.key_selector.db_connection", mock_db):
# Run multiple concurrent requests
tasks = [selector.get_next_key(1) for _ in range(5)]
results = await asyncio.gather(*tasks)
# All should return valid keys
assert all(r is not None for r in results)
@pytest.mark.asyncio
async def test_different_providers_independent(self, mock_db):
"""Test that different providers have independent indices."""
selector = KeySelector()
with patch("app.services.key_selector.db_connection", mock_db):
key1 = await selector.get_next_key(1)
key2 = await selector.get_next_key(2)
# Both should return valid keys from different providers
assert key1 is not None
assert key2 is not None
assert key1.provider_id == 1
assert key2.provider_id == 2