147 lines
5.8 KiB
Python
147 lines
5.8 KiB
Python
import pytest
|
||
from unittest.mock import AsyncMock, patch, MagicMock
|
||
from uuid import uuid4
|
||
import aiohttp
|
||
from tg_bot.application.services.rag_service import RAGService
|
||
from tests.metrics.test_hit_at_5 import calculate_hit_at_5, calculate_average_hit_at_5
|
||
|
||
|
||
class TestRAGIntegration:
|
||
|
||
@pytest.fixture
|
||
def rag_service(self):
|
||
service = RAGService()
|
||
from tg_bot.infrastructure.external.deepseek_client import DeepSeekClient
|
||
service.deepseek_client = DeepSeekClient()
|
||
return service
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_rag_service_search_documents_real_implementation(self, rag_service):
|
||
user_telegram_id = "123456789"
|
||
query = "трудовой договор"
|
||
user_uuid = str(uuid4())
|
||
collection_id = str(uuid4())
|
||
|
||
mock_user_response = AsyncMock()
|
||
mock_user_response.status = 200
|
||
mock_user_response.json = AsyncMock(return_value={"user_id": user_uuid})
|
||
|
||
mock_collections_response = AsyncMock()
|
||
mock_collections_response.status = 200
|
||
mock_collections_response.json = AsyncMock(return_value=[
|
||
{"collection_id": collection_id, "name": "Законы"}
|
||
])
|
||
|
||
mock_documents_response = AsyncMock()
|
||
mock_documents_response.status = 200
|
||
mock_documents_response.json = AsyncMock(return_value=[
|
||
{
|
||
"document_id": str(uuid4()),
|
||
"title": "Трудовой кодекс",
|
||
"content": "Содержание о трудовых договорах"
|
||
}
|
||
])
|
||
|
||
with patch('aiohttp.ClientSession') as mock_session_class:
|
||
mock_session = AsyncMock()
|
||
mock_session.__aenter__ = AsyncMock(return_value=mock_session)
|
||
mock_session.__aexit__ = AsyncMock(return_value=None)
|
||
mock_session.get = AsyncMock(side_effect=[
|
||
mock_user_response,
|
||
mock_collections_response,
|
||
mock_documents_response
|
||
])
|
||
mock_session_class.return_value = mock_session
|
||
|
||
result = await rag_service.search_documents_in_collections(user_telegram_id, query)
|
||
|
||
assert isinstance(result, list)
|
||
if result:
|
||
assert "document_id" in result[0]
|
||
assert "title" in result[0]
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_rag_service_generate_answer_real_implementation(self, rag_service):
|
||
question = "Какие права имеет работник?"
|
||
user_telegram_id = "123456789"
|
||
|
||
mock_documents = [
|
||
{
|
||
"document_id": str(uuid4()),
|
||
"title": "Трудовой кодекс",
|
||
"content": "Работник имеет право на...",
|
||
"collection_name": "Законы"
|
||
}
|
||
for _ in range(3)
|
||
]
|
||
|
||
with patch.object(rag_service, 'search_documents_in_collections') as mock_search, \
|
||
patch.object(rag_service.deepseek_client, 'chat_completion') as mock_deepseek:
|
||
|
||
mock_search.return_value = mock_documents
|
||
mock_deepseek.return_value = {
|
||
"content": "Работник имеет следующие права...",
|
||
"usage": {"prompt_tokens": 100, "completion_tokens": 200, "total_tokens": 300}
|
||
}
|
||
|
||
result = await rag_service.generate_answer_with_rag(question, user_telegram_id)
|
||
|
||
assert "answer" in result
|
||
assert "sources" in result
|
||
assert "usage" in result
|
||
assert len(result["sources"]) <= 5
|
||
assert result["answer"] != ""
|
||
|
||
for source in result["sources"]:
|
||
assert "title" in source
|
||
assert "collection" in source
|
||
assert "document_id" in source
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_rag_service_limits_to_top5_documents(self, rag_service):
|
||
question = "Тестовый вопрос"
|
||
user_telegram_id = "123456789"
|
||
|
||
many_documents = [
|
||
{
|
||
"document_id": str(uuid4()),
|
||
"title": f"Документ {i}",
|
||
"content": f"Содержание {i}",
|
||
"collection_name": "Коллекция"
|
||
}
|
||
for i in range(10)
|
||
]
|
||
|
||
with patch.object(rag_service, 'search_documents_in_collections') as mock_search, \
|
||
patch.object(rag_service.deepseek_client, 'chat_completion') as mock_deepseek:
|
||
|
||
mock_search.return_value = many_documents
|
||
mock_deepseek.return_value = {
|
||
"content": "Ответ",
|
||
"usage": {}
|
||
}
|
||
|
||
result = await rag_service.generate_answer_with_rag(question, user_telegram_id)
|
||
|
||
assert len(result["sources"]) == 5
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_rag_service_handles_empty_search_results(self, rag_service):
|
||
question = "Вопрос без документов"
|
||
user_telegram_id = "123456789"
|
||
|
||
with patch.object(rag_service, 'search_documents_in_collections') as mock_search, \
|
||
patch.object(rag_service.deepseek_client, 'chat_completion') as mock_deepseek:
|
||
|
||
mock_search.return_value = []
|
||
mock_deepseek.return_value = {
|
||
"content": "Ответ",
|
||
"usage": {}
|
||
}
|
||
|
||
result = await rag_service.generate_answer_with_rag(question, user_telegram_id)
|
||
|
||
assert result["sources"] == []
|
||
assert "Релевантные документы не найдены" in result.get("answer", "") or \
|
||
result["answer"] == "No relevant documents found"
|