355 lines
13 KiB
Python
355 lines
13 KiB
Python
import pytest
|
|
from unittest.mock import AsyncMock, patch, MagicMock
|
|
import httpx
|
|
|
|
from app.services.key_verifier import KeyStatus
|
|
from app.services.key_verifier import (
|
|
KeyVerifier,
|
|
KeyVerifierFactory,
|
|
ChatGPTKeyVerifier,
|
|
DeepSeekKeyVerifier,
|
|
KimiKeyVerifier,
|
|
QwenKeyVerifier,
|
|
PerplexityKeyVerifier,
|
|
GeminiKeyVerifier,
|
|
WenxinKeyVerifier,
|
|
DoubaoKeyVerifier,
|
|
YuanbaoKeyVerifier,
|
|
DefaultKeyVerifier,
|
|
)
|
|
|
|
|
|
def create_mock_client(mock_response: MagicMock) -> AsyncMock:
|
|
mock_client = AsyncMock()
|
|
mock_client.__aenter__ = AsyncMock(return_value=mock_client)
|
|
mock_client.__aexit__ = AsyncMock(return_value=None)
|
|
mock_client.post = AsyncMock(return_value=mock_response)
|
|
mock_client.get = AsyncMock(return_value=mock_response)
|
|
return mock_client
|
|
|
|
|
|
class TestKeyVerifierFactory:
|
|
def test_get_chatgpt_verifier(self):
|
|
verifier = KeyVerifierFactory.get_verifier("chatgpt")
|
|
assert isinstance(verifier, ChatGPTKeyVerifier)
|
|
|
|
def test_get_deepseek_verifier(self):
|
|
verifier = KeyVerifierFactory.get_verifier("deepseek")
|
|
assert isinstance(verifier, DeepSeekKeyVerifier)
|
|
|
|
def test_get_kimi_verifier(self):
|
|
verifier = KeyVerifierFactory.get_verifier("kimi")
|
|
assert isinstance(verifier, KimiKeyVerifier)
|
|
|
|
def test_get_qwen_verifier(self):
|
|
verifier = KeyVerifierFactory.get_verifier("qwen")
|
|
assert isinstance(verifier, QwenKeyVerifier)
|
|
|
|
def test_get_perplexity_verifier(self):
|
|
verifier = KeyVerifierFactory.get_verifier("perplexity")
|
|
assert isinstance(verifier, PerplexityKeyVerifier)
|
|
|
|
def test_get_gemini_verifier(self):
|
|
verifier = KeyVerifierFactory.get_verifier("gemini")
|
|
assert isinstance(verifier, GeminiKeyVerifier)
|
|
|
|
def test_get_unknown_verifier_returns_default(self):
|
|
verifier = KeyVerifierFactory.get_verifier("unknown_engine")
|
|
assert isinstance(verifier, DefaultKeyVerifier)
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_verify_delegates_to_correct_verifier(self):
|
|
with patch.object(ChatGPTKeyVerifier, 'verify', new_callable=AsyncMock) as mock_verify:
|
|
mock_verify.return_value = KeyStatus.ACTIVE
|
|
status = await KeyVerifierFactory.verify("chatgpt", "sk-test-key-12345")
|
|
mock_verify.assert_called_once_with("sk-test-key-12345")
|
|
assert status == KeyStatus.ACTIVE
|
|
|
|
|
|
class TestDefaultKeyVerifier:
|
|
@pytest.mark.asyncio
|
|
async def test_default_verifier_returns_unknown(self):
|
|
verifier = DefaultKeyVerifier()
|
|
status = await verifier.verify("any-key")
|
|
assert status == KeyStatus.UNKNOWN
|
|
|
|
|
|
class TestChatGPTKeyVerifier:
|
|
@pytest.mark.asyncio
|
|
async def test_verify_active_key_returns_active(self):
|
|
verifier = ChatGPTKeyVerifier()
|
|
mock_response = MagicMock()
|
|
mock_response.status_code = 200
|
|
|
|
mock_client = create_mock_client(mock_response)
|
|
|
|
with patch('httpx.AsyncClient', return_value=mock_client):
|
|
status = await verifier.verify("sk-test-key-12345")
|
|
assert status == KeyStatus.ACTIVE
|
|
mock_client.post.assert_called_once()
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_verify_invalid_key_returns_invalid(self):
|
|
verifier = ChatGPTKeyVerifier()
|
|
mock_response = MagicMock()
|
|
mock_response.status_code = 401
|
|
mock_response.text = "Invalid API key"
|
|
|
|
mock_client = create_mock_client(mock_response)
|
|
|
|
with patch('httpx.AsyncClient', return_value=mock_client):
|
|
status = await verifier.verify("sk-invalid-key")
|
|
assert status == KeyStatus.INVALID
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_verify_rate_limited_key_returns_rate_limited(self):
|
|
verifier = ChatGPTKeyVerifier()
|
|
mock_response = MagicMock()
|
|
mock_response.status_code = 429
|
|
mock_response.text = "Rate limit exceeded"
|
|
|
|
mock_client = create_mock_client(mock_response)
|
|
|
|
with patch('httpx.AsyncClient', return_value=mock_client):
|
|
status = await verifier.verify("sk-rate-limited-key")
|
|
assert status == KeyStatus.RATE_LIMITED
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_verify_network_error_returns_unknown(self):
|
|
verifier = ChatGPTKeyVerifier()
|
|
|
|
mock_client = create_mock_client(MagicMock())
|
|
mock_client.post = AsyncMock(side_effect=httpx.ConnectError("Connection failed"))
|
|
|
|
with patch('httpx.AsyncClient', return_value=mock_client):
|
|
status = await verifier.verify("sk-test-key")
|
|
assert status == KeyStatus.UNKNOWN
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_verify_timeout_returns_unknown(self):
|
|
verifier = ChatGPTKeyVerifier()
|
|
|
|
mock_client = create_mock_client(MagicMock())
|
|
mock_client.post = AsyncMock(side_effect=httpx.TimeoutException("Request timeout"))
|
|
|
|
with patch('httpx.AsyncClient', return_value=mock_client):
|
|
status = await verifier.verify("sk-test-key")
|
|
assert status == KeyStatus.UNKNOWN
|
|
|
|
|
|
class TestDeepSeekKeyVerifier:
|
|
@pytest.mark.asyncio
|
|
async def test_verify_active_key_returns_active(self):
|
|
verifier = DeepSeekKeyVerifier()
|
|
mock_response = MagicMock()
|
|
mock_response.status_code = 200
|
|
|
|
mock_client = create_mock_client(mock_response)
|
|
|
|
with patch('httpx.AsyncClient', return_value=mock_client):
|
|
status = await verifier.verify("sk-test-key-12345")
|
|
assert status == KeyStatus.ACTIVE
|
|
mock_client.post.assert_called_once()
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_verify_invalid_key_returns_invalid(self):
|
|
verifier = DeepSeekKeyVerifier()
|
|
mock_response = MagicMock()
|
|
mock_response.status_code = 401
|
|
|
|
mock_client = create_mock_client(mock_response)
|
|
|
|
with patch('httpx.AsyncClient', return_value=mock_client):
|
|
status = await verifier.verify("sk-invalid-key")
|
|
assert status == KeyStatus.INVALID
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_verify_rate_limited_key_returns_rate_limited(self):
|
|
verifier = DeepSeekKeyVerifier()
|
|
mock_response = MagicMock()
|
|
mock_response.status_code = 429
|
|
|
|
mock_client = create_mock_client(mock_response)
|
|
|
|
with patch('httpx.AsyncClient', return_value=mock_client):
|
|
status = await verifier.verify("sk-rate-limited-key")
|
|
assert status == KeyStatus.RATE_LIMITED
|
|
|
|
|
|
class TestKimiKeyVerifier:
|
|
@pytest.mark.asyncio
|
|
async def test_verify_active_key_returns_active(self):
|
|
verifier = KimiKeyVerifier()
|
|
mock_response = MagicMock()
|
|
mock_response.status_code = 200
|
|
|
|
mock_client = create_mock_client(mock_response)
|
|
|
|
with patch('httpx.AsyncClient', return_value=mock_client):
|
|
status = await verifier.verify("sk-test-key-12345")
|
|
assert status == KeyStatus.ACTIVE
|
|
mock_client.post.assert_called_once()
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_verify_invalid_key_returns_invalid(self):
|
|
verifier = KimiKeyVerifier()
|
|
mock_response = MagicMock()
|
|
mock_response.status_code = 401
|
|
|
|
mock_client = create_mock_client(mock_response)
|
|
|
|
with patch('httpx.AsyncClient', return_value=mock_client):
|
|
status = await verifier.verify("sk-invalid-key")
|
|
assert status == KeyStatus.INVALID
|
|
|
|
|
|
class TestQwenKeyVerifier:
|
|
@pytest.mark.asyncio
|
|
async def test_verify_active_key_returns_active(self):
|
|
verifier = QwenKeyVerifier()
|
|
mock_response = MagicMock()
|
|
mock_response.status_code = 200
|
|
|
|
mock_client = create_mock_client(mock_response)
|
|
|
|
with patch('httpx.AsyncClient', return_value=mock_client):
|
|
status = await verifier.verify("sk-test-key-12345")
|
|
assert status == KeyStatus.ACTIVE
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_verify_invalid_key_returns_invalid(self):
|
|
verifier = QwenKeyVerifier()
|
|
mock_response = MagicMock()
|
|
mock_response.status_code = 401
|
|
|
|
mock_client = create_mock_client(mock_response)
|
|
|
|
with patch('httpx.AsyncClient', return_value=mock_client):
|
|
status = await verifier.verify("sk-invalid-key")
|
|
assert status == KeyStatus.INVALID
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_verify_forbidden_key_returns_expired(self):
|
|
verifier = QwenKeyVerifier()
|
|
mock_response = MagicMock()
|
|
mock_response.status_code = 403
|
|
|
|
mock_client = create_mock_client(mock_response)
|
|
|
|
with patch('httpx.AsyncClient', return_value=mock_client):
|
|
status = await verifier.verify("sk-expired-key")
|
|
assert status == KeyStatus.EXPIRED
|
|
|
|
|
|
class TestPerplexityKeyVerifier:
|
|
@pytest.mark.asyncio
|
|
async def test_verify_active_key_returns_active(self):
|
|
verifier = PerplexityKeyVerifier()
|
|
mock_response = MagicMock()
|
|
mock_response.status_code = 200
|
|
|
|
mock_client = create_mock_client(mock_response)
|
|
|
|
with patch('httpx.AsyncClient', return_value=mock_client):
|
|
status = await verifier.verify("sk-test-key-12345")
|
|
assert status == KeyStatus.ACTIVE
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_verify_invalid_key_returns_invalid(self):
|
|
verifier = PerplexityKeyVerifier()
|
|
mock_response = MagicMock()
|
|
mock_response.status_code = 401
|
|
|
|
mock_client = create_mock_client(mock_response)
|
|
|
|
with patch('httpx.AsyncClient', return_value=mock_client):
|
|
status = await verifier.verify("sk-invalid-key")
|
|
assert status == KeyStatus.INVALID
|
|
|
|
|
|
class TestGeminiKeyVerifier:
|
|
@pytest.mark.asyncio
|
|
async def test_verify_active_key_returns_active(self):
|
|
verifier = GeminiKeyVerifier()
|
|
mock_response = MagicMock()
|
|
mock_response.status_code = 200
|
|
|
|
mock_client = create_mock_client(mock_response)
|
|
|
|
with patch('httpx.AsyncClient', return_value=mock_client):
|
|
status = await verifier.verify("sk-test-key-12345")
|
|
assert status == KeyStatus.ACTIVE
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_verify_invalid_key_returns_invalid(self):
|
|
verifier = GeminiKeyVerifier()
|
|
mock_response = MagicMock()
|
|
mock_response.status_code = 403
|
|
|
|
mock_client = create_mock_client(mock_response)
|
|
|
|
with patch('httpx.AsyncClient', return_value=mock_client):
|
|
status = await verifier.verify("sk-invalid-key")
|
|
assert status == KeyStatus.INVALID
|
|
|
|
|
|
class TestOtherEngineVerifiers:
|
|
@pytest.mark.asyncio
|
|
async def test_wenxin_verifier_returns_unknown(self):
|
|
verifier = WenxinKeyVerifier()
|
|
status = await verifier.verify("test-key")
|
|
assert status == KeyStatus.UNKNOWN
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_doubao_verifier_returns_unknown(self):
|
|
verifier = DoubaoKeyVerifier()
|
|
status = await verifier.verify("test-key")
|
|
assert status == KeyStatus.UNKNOWN
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_yuanbao_verifier_returns_unknown(self):
|
|
verifier = YuanbaoKeyVerifier()
|
|
status = await verifier.verify("test-key")
|
|
assert status == KeyStatus.UNKNOWN
|
|
|
|
|
|
class TestKeyVerifierIntegrationWithAPIManager:
|
|
@pytest.mark.asyncio
|
|
async def test_api_key_manager_verify_uses_factory(self):
|
|
from app.services.api_key_manager import APIKeyManager
|
|
|
|
manager = APIKeyManager()
|
|
|
|
with patch.object(KeyVerifierFactory, 'verify', new_callable=AsyncMock) as mock_verify:
|
|
mock_verify.return_value = KeyStatus.ACTIVE
|
|
status = await manager.verify_key("chatgpt", "sk-test-key-1234567890")
|
|
assert status == KeyStatus.ACTIVE
|
|
mock_verify.assert_called_once_with("chatgpt", "sk-test-key-1234567890")
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_api_key_manager_verify_handles_factory_error(self):
|
|
from app.services.api_key_manager import APIKeyManager
|
|
|
|
manager = APIKeyManager()
|
|
|
|
with patch.object(KeyVerifierFactory, 'verify', new_callable=AsyncMock) as mock_verify:
|
|
mock_verify.side_effect = Exception("Network error")
|
|
status = await manager.verify_key("chatgpt", "sk-test-key-1234567890")
|
|
assert status == KeyStatus.UNKNOWN
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_api_key_manager_verify_short_key_returns_invalid(self):
|
|
from app.services.api_key_manager import APIKeyManager
|
|
|
|
manager = APIKeyManager()
|
|
status = await manager.verify_key("chatgpt", "short")
|
|
assert status == KeyStatus.INVALID
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_api_key_manager_verify_empty_key_returns_invalid(self):
|
|
from app.services.api_key_manager import APIKeyManager
|
|
|
|
manager = APIKeyManager()
|
|
status = await manager.verify_key("chatgpt", "")
|
|
assert status == KeyStatus.INVALID
|