tests

tests/integration/test_rag_integration.py

@@ -0,0 +1,146 @@
+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"