tests/test_document_processor.py - main - cppv/hackathon13

tests/test_document_processor.py
178 строк7 KB
cppv
hackathon13
25 май 2025, 11:54
25 май 2025, 11:548af8c2d
import unittest
from unittest.mock import patch, MagicMock
import sys
import os

# добавляем путь чтоб работало
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))

from app.document_processor import DocumentProcessor

class TestDocumentProcessor(unittest.TestCase):
    """
    тесты для документов
    """
    
    def setUp(self):
        """
        подготовка к тестам
        """
        self.docs_dir = "/tmp/test_docs"
        self.index_path = "/tmp/test_index"
        
        # создаем папки
        os.makedirs(self.docs_dir, exist_ok=True)
        os.makedirs(os.path.dirname(self.index_path), exist_ok=True)
        
        # делаем тестовый файл
        with open(os.path.join(self.docs_dir, "test.md"), "w") as f:
            f.write("# Test Document\n\nThis is a test document for LangChain.")
        
        self.processor = DocumentProcessor(self.docs_dir, self.index_path)
        
    def tearDown(self):
        """
        убираем за собой
        """
        # чистим все что создали
        if os.path.exists(os.path.join(self.docs_dir, "test.md")):
            os.remove(os.path.join(self.docs_dir, "test.md"))
        
        if os.path.exists(self.docs_dir):
            os.rmdir(self.docs_dir)
        
        if os.path.exists(self.index_path):
            import shutil
            shutil.rmtree(os.path.dirname(self.index_path), ignore_errors=True)
    
    def test_load_documents(self):
        """
        проверяем загрузку файлов
        """
        documents = self.processor.load_documents()
        
        # должно что-то загрузиться
        self.assertGreater(len(documents), 0)
        
        # проверяем что там внутри норм текст
        self.assertIn("Test Document", documents[0].page_content)
    
    def test_split_documents(self):
        """
        проверяем что документы режутся на куски
        """
        documents = self.processor.load_documents()
        chunks = self.processor.split_documents(documents, chunk_size=10, chunk_overlap=2)
        
        # должно получиться больше кусков чем документов
        self.assertGreater(len(chunks), len(documents))
    
    @patch('app.document_processor.FAISS')
    def test_create_index(self, mock_faiss):
        """
        проверяем создание индекса
        """
        # готовим моки
        mock_vectorstore = MagicMock()
        mock_faiss.from_documents.return_value = mock_vectorstore
        
        documents = self.processor.load_documents()
        chunks = self.processor.split_documents(documents)
        mock_embeddings = MagicMock()
        
        vectorstore = self.processor.create_index(chunks, mock_embeddings)
        
        # проверяем что вызвалось как надо
        mock_faiss.from_documents.assert_called_once_with(chunks, mock_embeddings)
        
        # и вернулось что надо
        self.assertEqual(vectorstore, mock_vectorstore)
    
    @patch('app.document_processor.FAISS')
    def test_save_and_load_index(self, mock_faiss):
        """
        проверяем сохранение и загрузку индекса
        """
        # мокаем что надо
        mock_vectorstore = MagicMock()
        mock_faiss.load_local.return_value = mock_vectorstore
        
        mock_embeddings = MagicMock()
        
        # пробуем загрузить
        vectorstore = self.processor.load_index(mock_embeddings)
        
        # не должно сработать т.к. нет файла
        mock_faiss.load_local.assert_not_called()
        
        # должен быть None
        self.assertIsNone(vectorstore)
        
        # теперь делаем вид что файл есть
        with patch('os.path.exists', return_value=True):
            vectorstore = self.processor.load_index(mock_embeddings)
            
            # должно загрузиться
            mock_faiss.load_local.assert_called_once_with(self.index_path, mock_embeddings)
            
            # и вернуть что надо
            self.assertEqual(vectorstore, mock_vectorstore)
    
    @patch('app.document_processor.DocumentProcessor.load_index')
    @patch('app.document_processor.DocumentProcessor.load_documents')
    @patch('app.document_processor.DocumentProcessor.split_documents')
    @patch('app.document_processor.DocumentProcessor.create_index')
    @patch('app.document_processor.DocumentProcessor.save_index')
    def test_process_documents(self, mock_save_index, mock_create_index, 
                              mock_split_documents, mock_load_documents, mock_load_index):
        """
        проверяем что все работает вместе
        """
        # готовим моки
        mock_vectorstore = MagicMock()
        mock_load_index.return_value = None  # нет индекса
        mock_documents = [MagicMock()]
        mock_load_documents.return_value = mock_documents
        mock_chunks = [MagicMock(), MagicMock()]
        mock_split_documents.return_value = mock_chunks
        mock_create_index.return_value = mock_vectorstore
        
        mock_embeddings = MagicMock()
        
        # запускаем
        result = self.processor.process_documents(mock_embeddings)
        
        # все должно вызваться
        mock_load_index.assert_called_once_with(mock_embeddings)
        mock_load_documents.assert_called_once()
        mock_split_documents.assert_called_once_with(mock_documents)
        mock_create_index.assert_called_once_with(mock_chunks, mock_embeddings)
        mock_save_index.assert_called_once_with(mock_vectorstore)
        
        # и вернуться что надо
        self.assertEqual(result, mock_vectorstore)
        
        # теперь пробуем когда индекс уже есть
        mock_load_index.return_value = mock_vectorstore
        
        # чистим моки
        mock_load_documents.reset_mock()
        mock_split_documents.reset_mock()
        mock_create_index.reset_mock()
        mock_save_index.reset_mock()
        
        # запускаем еще раз
        result = self.processor.process_documents(mock_embeddings)
        
        # должен только загрузиться индекс
        mock_load_index.assert_called_with(mock_embeddings)
        mock_load_documents.assert_not_called()
        mock_split_documents.assert_not_called()
        mock_create_index.assert_not_called()
        mock_save_index.assert_not_called()
        
        # и вернуться тот же индекс
        self.assertEqual(result, mock_vectorstore)

if __name__ == '__main__':
    unittest.main()