93 lines
2.9 KiB
Python
93 lines
2.9 KiB
Python
"""Тесты для GigaChat клиента."""
|
||
import pytest
|
||
from unittest.mock import AsyncMock, MagicMock, patch
|
||
|
||
from agents.gigachat_client import GigaChatClient
|
||
from models.gigachat_types import GigaChatMessage, GigaChatResponse, GigaChatUsage, GigaChatChoice
|
||
from services.token_manager import TokenManager
|
||
|
||
|
||
@pytest.fixture
|
||
def token_manager():
|
||
"""Фикстура для TokenManager."""
|
||
manager = TokenManager(
|
||
client_id="test_id",
|
||
client_secret="test_secret",
|
||
)
|
||
manager._access_token = "test_token"
|
||
manager._expires_at = 9999999999
|
||
return manager
|
||
|
||
|
||
@pytest.fixture
|
||
def gigachat_client(token_manager):
|
||
"""Фикстура для GigaChatClient."""
|
||
return GigaChatClient(token_manager=token_manager)
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_chat_success(gigachat_client):
|
||
"""Тест успешного запроса к GigaChat."""
|
||
mock_response = GigaChatResponse(
|
||
id="test_id",
|
||
object="chat.completion",
|
||
created=1234567890,
|
||
model="GigaChat-2",
|
||
choices=[
|
||
GigaChatChoice(
|
||
message=GigaChatMessage(role="assistant", content="Тестовый ответ"),
|
||
index=0,
|
||
finish_reason="stop",
|
||
)
|
||
],
|
||
usage=GigaChatUsage(
|
||
prompt_tokens=10,
|
||
completion_tokens=5,
|
||
total_tokens=15,
|
||
),
|
||
)
|
||
|
||
with patch("aiohttp.ClientSession.post") as mock_post:
|
||
mock_response_obj = AsyncMock()
|
||
mock_response_obj.status = 200
|
||
mock_response_obj.json = AsyncMock(return_value=mock_response.model_dump())
|
||
mock_post.return_value.__aenter__.return_value = mock_response_obj
|
||
|
||
response = await gigachat_client.chat("Привет!")
|
||
|
||
assert response == "Тестовый ответ"
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_chat_with_context(gigachat_client):
|
||
"""Тест запроса с контекстом."""
|
||
context = [
|
||
GigaChatMessage(role="system", content="Ты помощник"),
|
||
GigaChatMessage(role="user", content="Привет"),
|
||
]
|
||
|
||
mock_response = GigaChatResponse(
|
||
id="test_id",
|
||
object="chat.completion",
|
||
created=1234567890,
|
||
model="GigaChat-2",
|
||
choices=[
|
||
GigaChatChoice(
|
||
message=GigaChatMessage(role="assistant", content="Ответ с контекстом"),
|
||
index=0,
|
||
)
|
||
],
|
||
usage=GigaChatUsage(prompt_tokens=20, completion_tokens=10, total_tokens=30),
|
||
)
|
||
|
||
with patch("aiohttp.ClientSession.post") as mock_post:
|
||
mock_response_obj = AsyncMock()
|
||
mock_response_obj.status = 200
|
||
mock_response_obj.json = AsyncMock(return_value=mock_response.model_dump())
|
||
mock_post.return_value.__aenter__.return_value = mock_response_obj
|
||
|
||
response = await gigachat_client.chat("Как дела?", context=context)
|
||
|
||
assert response == "Ответ с контекстом"
|
||
|