forked from HSE_team/BetterCallPraskovia
115 lines
4.6 KiB
Python
115 lines
4.6 KiB
Python
import pytest
|
|
from unittest.mock import AsyncMock, patch, MagicMock
|
|
import httpx
|
|
from tg_bot.infrastructure.external.deepseek_client import DeepSeekClient
|
|
from tg_bot.infrastructure.external.deepseek_client import DeepSeekAPIError
|
|
|
|
|
|
class TestDeepSeekClient:
|
|
|
|
@pytest.fixture
|
|
def deepseek_client(self):
|
|
return DeepSeekClient(api_key="test_key", api_url="https://api.test.com/v1/chat/completions")
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_chat_completion_success(self, deepseek_client):
|
|
messages = [
|
|
{"role": "user", "content": "Тестовый вопрос"}
|
|
]
|
|
|
|
mock_response_data = {
|
|
"choices": [{
|
|
"message": {
|
|
"content": "Тестовый ответ от DeepSeek"
|
|
}
|
|
}],
|
|
"usage": {
|
|
"prompt_tokens": 10,
|
|
"completion_tokens": 20,
|
|
"total_tokens": 30
|
|
}
|
|
}
|
|
|
|
with patch('httpx.AsyncClient') as mock_client:
|
|
mock_response = AsyncMock()
|
|
mock_response.status_code = 200
|
|
mock_response.json.return_value = mock_response_data
|
|
mock_response.raise_for_status = MagicMock()
|
|
|
|
mock_client_instance = AsyncMock()
|
|
mock_client_instance.__aenter__ = AsyncMock(return_value=mock_client_instance)
|
|
mock_client_instance.__aexit__ = AsyncMock(return_value=None)
|
|
mock_client_instance.post = AsyncMock(return_value=mock_response)
|
|
mock_client.return_value = mock_client_instance
|
|
|
|
result = await deepseek_client.chat_completion(messages)
|
|
|
|
assert "content" in result
|
|
assert result["content"] == "Тестовый ответ от DeepSeek"
|
|
assert "usage" in result
|
|
assert result["usage"]["total_tokens"] == 30
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_chat_completion_no_api_key(self):
|
|
client = DeepSeekClient(api_key=None)
|
|
messages = [{"role": "user", "content": "Вопрос"}]
|
|
|
|
result = await client.chat_completion(messages)
|
|
|
|
assert "content" in result
|
|
assert "DEEPSEEK_API_KEY" in result["content"] or "не установлен" in result["content"]
|
|
assert result["usage"]["total_tokens"] == 0
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_chat_completion_api_error(self, deepseek_client):
|
|
import httpx
|
|
messages = [{"role": "user", "content": "Вопрос"}]
|
|
|
|
with patch('httpx.AsyncClient') as mock_client:
|
|
mock_response = AsyncMock()
|
|
mock_response.status_code = 401
|
|
mock_response.raise_for_status.side_effect = httpx.HTTPStatusError(
|
|
"Unauthorized", request=MagicMock(), response=mock_response
|
|
)
|
|
|
|
mock_client_instance = AsyncMock()
|
|
mock_client_instance.__aenter__ = AsyncMock(return_value=mock_client_instance)
|
|
mock_client_instance.__aexit__ = AsyncMock(return_value=None)
|
|
mock_client_instance.post = AsyncMock(return_value=mock_response)
|
|
mock_client.return_value = mock_client_instance
|
|
|
|
with pytest.raises(DeepSeekAPIError):
|
|
await deepseek_client.chat_completion(messages)
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_chat_completion_with_parameters(self, deepseek_client):
|
|
messages = [{"role": "user", "content": "Вопрос"}]
|
|
|
|
mock_response_data = {
|
|
"choices": [{"message": {"content": "Ответ"}}],
|
|
"usage": {"prompt_tokens": 5, "completion_tokens": 10, "total_tokens": 15}
|
|
}
|
|
|
|
with patch('httpx.AsyncClient') as mock_client:
|
|
mock_response = AsyncMock()
|
|
mock_response.status_code = 200
|
|
mock_response.json.return_value = mock_response_data
|
|
mock_response.raise_for_status = MagicMock()
|
|
|
|
mock_client_instance = AsyncMock()
|
|
mock_client_instance.__aenter__ = AsyncMock(return_value=mock_client_instance)
|
|
mock_client_instance.__aexit__ = AsyncMock(return_value=None)
|
|
mock_client_instance.post = AsyncMock(return_value=mock_response)
|
|
mock_client.return_value = mock_client_instance
|
|
|
|
result = await deepseek_client.chat_completion(
|
|
messages,
|
|
model="deepseek-chat",
|
|
temperature=0.7,
|
|
max_tokens=100
|
|
)
|
|
|
|
assert result["content"] == "Ответ"
|
|
call_args = mock_client_instance.post.call_args
|
|
assert call_args is not None
|