tests/test_document_handler.py - master - alsoalgo/CTF

tests/test_document_handler.py
452 строки22 KB
alsoalgo
some improvements
14 дек 2025, 22:23
14 дек 2025, 22:2337d2ec0
"""
Тесты для DocumentHandler
"""
import pytest
import pandas as pd
import os
import tempfile
import logging
from pathlib import Path
from unittest.mock import Mock, patch, MagicMock

from src.components.document_handler import DocumentHandler


class TestDocumentHandler:
    """Тесты для класса DocumentHandler"""
    
    def setup_method(self):
        """Настройка перед каждым тестом"""
        self.handler = DocumentHandler()
        self.test_dir = Path(__file__).parent.parent
    
    def test_load_csv_document(self):
        """Тест загрузки CSV файла"""
        csv_path = self.test_dir / "examples" / "sample_invoice.csv"
        if csv_path.exists():
            df = self.handler.load_csv_document(str(csv_path))
            assert df is not None
            assert isinstance(df, pd.DataFrame)
            assert not df.empty
            assert "Invoice Number" in df.columns or "Supplier" in df.columns
    
    def test_load_nonexistent_file(self):
        """Тест загрузки несуществующего файла"""
        df = self.handler.load_document("/nonexistent/file.csv")
        assert df is None
    
    def test_load_xlsx_document(self):
        """Тест загрузки XLSX файла"""
        xlsx_path = self.test_dir / "examples" / "recap. Georgian Solutions.xlsx"
        if xlsx_path.exists():
            df = self.handler.load_xlsx_document(str(xlsx_path))
            assert df is not None
            assert isinstance(df, pd.DataFrame)
            assert not df.empty
    
    def test_save_dataframe_to_csv(self):
        """Тест сохранения DataFrame в CSV"""
        # Создаем тестовый DataFrame
        test_df = pd.DataFrame({
            'Field': ['Test1', 'Test2'],
            'Value': ['Value1', 'Value2']
        })
        
        # Сохраняем во временный файл
        with tempfile.NamedTemporaryFile(mode='w', suffix='.csv', delete=False) as tmp:
            tmp_path = tmp.name
        
        try:
            self.handler.save_dataframe_to_csv(test_df, tmp_path)
            assert os.path.exists(tmp_path)
            
            # Проверяем, что файл можно загрузить обратно
            loaded_df = pd.read_csv(tmp_path)
            assert len(loaded_df) == 2
            assert 'Field' in loaded_df.columns
        finally:
            if os.path.exists(tmp_path):
                os.remove(tmp_path)
    
    def test_load_document_csv(self):
        """Тест универсального метода load_document для CSV"""
        csv_path = self.test_dir / "examples" / "sample_invoice.csv"
        if csv_path.exists():
            df = self.handler.load_document(str(csv_path))
            assert df is not None
            assert isinstance(df, pd.DataFrame)
    
    def test_load_document_xlsx(self):
        """Тест универсального метода load_document для XLSX"""
        xlsx_path = self.test_dir / "examples" / "recap. Georgian Solutions.xlsx"
        if xlsx_path.exists():
            df = self.handler.load_document(str(xlsx_path))
            assert df is not None
            assert isinstance(df, pd.DataFrame)
    
    def test_load_document_unsupported_format(self):
        """Тест загрузки неподдерживаемого формата"""
        # Создаем временный файл с неподдерживаемым расширением
        with tempfile.NamedTemporaryFile(mode='w', suffix='.txt', delete=False) as tmp:
            tmp.write("test content")
            tmp_path = tmp.name
        
        try:
            df = self.handler.load_document(tmp_path)
            assert df is None
        finally:
            if os.path.exists(tmp_path):
                os.remove(tmp_path)
    
    def test_load_document_with_track_sources(self):
        """Тест загрузки документа с отслеживанием источников"""
        csv_path = self.test_dir / "examples" / "sample_invoice.csv"
        if csv_path.exists():
            result = self.handler.load_document(str(csv_path), track_sources=True)
            # Должен вернуть tuple (DataFrame, sources_info)
            assert isinstance(result, tuple)
            assert len(result) == 2
            df, sources_info = result
            assert df is not None
            assert isinstance(df, pd.DataFrame)
            assert sources_info is not None
            assert isinstance(sources_info, dict)
    
    def test_load_document_without_track_sources(self):
        """Тест загрузки документа без отслеживания источников"""
        csv_path = self.test_dir / "examples" / "sample_invoice.csv"
        if csv_path.exists():
            result = self.handler.load_document(str(csv_path), track_sources=False)
            # Может вернуть DataFrame или tuple для обратной совместимости
            if isinstance(result, tuple):
                df = result[0]
            else:
                df = result
            assert df is not None
            assert isinstance(df, pd.DataFrame)
    
    def test_load_pdf_comprehensive_analysis(self):
        """Тест комплексного анализа PDF (текст + таблицы)"""
        pdf_path = self.test_dir / "examples" / "forwarder invoice sample.pdf"
        if pdf_path.exists():
            result = self.handler.load_pdf_document(str(pdf_path), track_sources=True)
            assert result is not None
            assert isinstance(result, tuple)
            df, sources_info = result
            assert df is not None
            assert isinstance(df, pd.DataFrame)
            assert sources_info is not None
            assert isinstance(sources_info, dict)
            # Проверяем структуру sources_info
            assert 'file' in sources_info
            assert 'text_sources' in sources_info or 'table_sources' in sources_info
    
    def test_load_pdf_with_text_extraction(self):
        """Тест извлечения текста из PDF"""
        pdf_path = self.test_dir / "examples" / "forwarder invoice sample.pdf"
        if pdf_path.exists():
            result = self.handler.load_pdf_document(str(pdf_path), track_sources=True)
            if result:
                df, sources_info = result
                if df is not None and not df.empty:
                    # Проверяем, что есть данные с Source='text'
                    if 'Source' in df.columns:
                        text_rows = df[df['Source'] == 'text']
                        assert len(text_rows) >= 0  # Может быть 0, если нет текста
    
    def test_load_pdf_with_table_extraction(self):
        """Тест извлечения таблиц из PDF"""
        pdf_path = self.test_dir / "examples" / "forwarder invoice sample.pdf"
        if pdf_path.exists():
            result = self.handler.load_pdf_document(str(pdf_path), track_sources=True)
            if result:
                df, sources_info = result
                if df is not None and not df.empty:
                    # Проверяем, что есть данные с Source='table'
                    if 'Source' in df.columns:
                        table_rows = df[df['Source'] == 'table']
                        assert len(table_rows) >= 0  # Может быть 0, если нет таблиц
                    # Проверяем информацию о таблицах в sources_info
                    if 'table_sources' in sources_info:
                        assert isinstance(sources_info['table_sources'], list)
    
    def test_load_pdf_sources_info_structure(self):
        """Тест структуры sources_info для PDF"""
        pdf_path = self.test_dir / "examples" / "forwarder invoice sample.pdf"
        if pdf_path.exists():
            result = self.handler.load_pdf_document(str(pdf_path), track_sources=True)
            if result:
                df, sources_info = result
                if sources_info:
                    # Проверяем обязательные поля
                    assert 'file' in sources_info
                    assert 'text_sources' in sources_info
                    assert 'table_sources' in sources_info
                    assert 'ocr_sources' in sources_info
                    assert 'images_found' in sources_info
                    assert isinstance(sources_info['text_sources'], list)
                    assert isinstance(sources_info['table_sources'], list)
                    assert isinstance(sources_info['ocr_sources'], list)
                    assert isinstance(sources_info['images_found'], bool)
    
    def test_load_pdf_aggregation(self):
        """Тест агрегации данных из разных источников"""
        pdf_path = self.test_dir / "examples" / "forwarder invoice sample.pdf"
        if pdf_path.exists():
            result = self.handler.load_pdf_document(str(pdf_path), track_sources=True)
            if result:
                df, sources_info = result
                if df is not None and not df.empty:
                    # Проверяем, что DataFrame содержит данные
                    assert len(df) > 0
                    # Проверяем наличие колонки Source
                    if 'Source' in df.columns:
                        # Должны быть данные хотя бы из одного источника
                        sources = df['Source'].unique()
                        assert len(sources) > 0
    
    @patch('src.components.document_handler.OCRHandler')
    def test_load_pdf_with_ocr_when_images_found(self, mock_ocr_class):
        """Тест использования OCR при обнаружении изображений"""
        # Мокаем OCR handler
        mock_ocr = MagicMock()
        mock_ocr_class.return_value = mock_ocr
        
        # Мокаем результат OCR
        mock_ocr.extract_structured_data.return_value = pd.DataFrame({
            'Field': ['Test'],
            'Value': ['Value']
        })
        
        handler = DocumentHandler(use_ocr=True, ocr_api_key="test_key")
        handler.ocr_handler = mock_ocr
        
        pdf_path = self.test_dir / "examples" / "forwarder invoice sample.pdf"
        if pdf_path.exists():
            result = handler.load_pdf_document(str(pdf_path), track_sources=True)
            # OCR может быть вызван, если обнаружены изображения
            # Проверяем, что метод существует и может быть вызван
            assert hasattr(handler, '_extract_with_ocr_comprehensive')
    
    def test_extract_with_ocr_comprehensive_no_ocr_handler(self):
        """Тест _extract_with_ocr_comprehensive без OCR handler"""
        handler = DocumentHandler(use_ocr=False)
        result = handler._extract_with_ocr_comprehensive("/fake/path.pdf", track_sources=False)
        assert result is None
    
    @patch('src.components.document_handler.OCRHandler')
    def test_extract_with_ocr_comprehensive_with_pdf_text(self, mock_ocr_class):
        """Тест _extract_with_ocr_comprehensive с pdf-text движком"""
        mock_ocr = MagicMock()
        mock_ocr_class.return_value = mock_ocr
        
        # Мокаем успешный результат
        test_df = pd.DataFrame({
            'Field': ['Supplier', 'Amount'],
            'Value': ['Test Supplier', '1000']
        })
        mock_ocr.extract_structured_data.return_value = test_df
        
        handler = DocumentHandler(use_ocr=True, ocr_api_key="test_key")
        handler.ocr_handler = mock_ocr
        
        with tempfile.NamedTemporaryFile(mode='wb', suffix='.pdf', delete=False) as tmp:
            tmp.write(b"fake pdf content")
            tmp_path = tmp.name
        
        try:
            result = handler._extract_with_ocr_comprehensive(tmp_path, track_sources=True)
            assert result is not None
            df, sources = result
            assert df is not None
            assert isinstance(df, pd.DataFrame)
            assert sources is not None
            assert sources.get('method') == 'ocr_pdf_text'
        finally:
            if os.path.exists(tmp_path):
                os.remove(tmp_path)
    
    @patch('src.components.document_handler.OCRHandler')
    def test_extract_with_ocr_comprehensive_with_mistral_ocr(self, mock_ocr_class):
        """Тест _extract_with_ocr_comprehensive с mistral-ocr движком"""
        mock_ocr = MagicMock()
        mock_ocr_class.return_value = mock_ocr
        
        # Мокаем, что pdf-text не сработал
        mock_ocr.extract_structured_data.return_value = None
        
        # Мокаем успешный результат mistral-ocr
        mock_ocr.extract_text_with_ocr.return_value = "Field | Value\nSupplier | Test Supplier"
        
        handler = DocumentHandler(use_ocr=True, ocr_api_key="test_key")
        handler.ocr_handler = mock_ocr
        
        with tempfile.NamedTemporaryFile(mode='wb', suffix='.pdf', delete=False) as tmp:
            tmp.write(b"fake pdf content")
            tmp_path = tmp.name
        
        try:
            result = handler._extract_with_ocr_comprehensive(tmp_path, track_sources=True)
            # Может вернуть результат или None в зависимости от парсинга
            if result:
                df, sources = result
                assert df is not None
                assert isinstance(df, pd.DataFrame)
        finally:
            if os.path.exists(tmp_path):
                os.remove(tmp_path)
    
    def test_load_pdf_comprehensive_without_ocr(self):
        """Тест комплексного анализа PDF без OCR"""
        handler = DocumentHandler(use_ocr=False)
        pdf_path = self.test_dir / "examples" / "forwarder invoice sample.pdf"
        if pdf_path.exists():
            result = handler.load_pdf_document(str(pdf_path), track_sources=True)
            # Должен работать даже без OCR (извлечение текста и таблиц)
            if result:
                df, sources_info = result
                # Может быть None если файл не читается, или DataFrame если читается
                assert df is None or isinstance(df, pd.DataFrame)
    
    def test_load_pdf_empty_file(self):
        """Тест загрузки пустого PDF файла"""
        with tempfile.NamedTemporaryFile(mode='wb', suffix='.pdf', delete=False) as tmp:
            # Создаем минимальный пустой PDF (не валидный, но для теста)
            tmp.write(b"not a valid pdf")
            tmp_path = tmp.name
        
        try:
            result = self.handler.load_pdf_document(tmp_path, track_sources=True)
            # Может вернуть None или tuple с None
            if result:
                if isinstance(result, tuple):
                    df, sources_info = result
                    assert df is None or isinstance(df, pd.DataFrame)
        except Exception:
            # Ожидаем исключение для невалидного PDF
            pass
        finally:
            if os.path.exists(tmp_path):
                os.remove(tmp_path)
    
    def test_load_pdf_comprehensive_returns_correct_format(self):
        """Тест формата возвращаемых данных при комплексном анализе"""
        pdf_path = self.test_dir / "examples" / "forwarder invoice sample.pdf"
        if pdf_path.exists():
            result = self.handler.load_pdf_document(str(pdf_path), track_sources=True)
            if result:
                assert isinstance(result, tuple)
                df, sources_info = result
                if df is not None:
                    assert isinstance(df, pd.DataFrame)
                    # Проверяем наличие колонки Source если есть данные
                    if len(df) > 0:
                        # Может быть колонка Source или Content
                        assert 'Source' in df.columns or 'Content' in df.columns or len(df.columns) > 0
    
    @patch('src.components.document_handler.OCRHandler')
    def test_extract_with_ocr_comprehensive_handles_api_errors(self, mock_ocr_class):
        """Тест обработки ошибок API при OCR извлечении"""
        mock_ocr = MagicMock()
        mock_ocr_class.return_value = mock_ocr
        
        # Мокаем ошибку API
        mock_ocr.extract_structured_data.return_value = None
        mock_ocr.extract_text_with_ocr.return_value = None
        
        handler = DocumentHandler(use_ocr=True, ocr_api_key="test_key")
        handler.ocr_handler = mock_ocr
        
        with tempfile.NamedTemporaryFile(mode='wb', suffix='.pdf', delete=False) as tmp:
            tmp.write(b"fake pdf content")
            tmp_path = tmp.name
        
        try:
            result = handler._extract_with_ocr_comprehensive(tmp_path, track_sources=True)
            # Должен вернуть None при ошибке
            assert result is None
        finally:
            if os.path.exists(tmp_path):
                os.remove(tmp_path)
    
    @patch('src.components.document_handler.OCRHandler')
    def test_load_pdf_continues_after_ocr_error(self, mock_ocr_class):
        """Тест продолжения работы после ошибки OCR"""
        mock_ocr = MagicMock()
        mock_ocr_class.return_value = mock_ocr
        
        # Мокаем ошибку OCR
        mock_ocr.extract_structured_data.return_value = None
        mock_ocr.extract_text_with_ocr.return_value = None
        
        handler = DocumentHandler(use_ocr=True, ocr_api_key="test_key")
        handler.ocr_handler = mock_ocr
        
        pdf_path = self.test_dir / "examples" / "forwarder invoice sample.pdf"
        if pdf_path.exists():
            # Должен работать даже если OCR не сработал (есть текст/таблицы)
            result = handler.load_pdf_document(str(pdf_path), track_sources=True)
            # Может вернуть данные из текста/таблиц даже если OCR не сработал
            if result:
                df, sources_info = result
                # Проверяем, что хотя бы структура правильная
                assert isinstance(result, tuple)
    
    def test_load_pdf_logging_integration(self, caplog):
        """Тест логирования при комплексном анализе PDF"""
        import logging
        logging.basicConfig(level=logging.INFO)
        
        pdf_path = self.test_dir / "examples" / "forwarder invoice sample.pdf"
        if pdf_path.exists():
            with caplog.at_level(logging.INFO):
                result = self.handler.load_pdf_document(str(pdf_path), track_sources=True)
                # Проверяем, что есть логи
                log_messages = [record.message for record in caplog.records]
                assert any("PDF содержит" in msg or "страниц" in msg for msg in log_messages)
    
    def test_extract_with_ocr_comprehensive_error_handling(self):
        """Тест обработки исключений в _extract_with_ocr_comprehensive"""
        handler = DocumentHandler(use_ocr=False)
        # Должен вернуть None если нет OCR handler
        result = handler._extract_with_ocr_comprehensive("/fake/path.pdf", track_sources=False)
        assert result is None
    
    @patch('src.components.document_handler.OCRHandler')
    def test_extract_with_ocr_comprehensive_exception_handling(self, mock_ocr_class):
        """Тест обработки исключений при OCR извлечении"""
        mock_ocr = MagicMock()
        mock_ocr_class.return_value = mock_ocr
        
        # Мокаем исключение
        mock_ocr.extract_structured_data.side_effect = Exception("Test error")
        
        handler = DocumentHandler(use_ocr=True, ocr_api_key="test_key")
        handler.ocr_handler = mock_ocr
        
        with tempfile.NamedTemporaryFile(mode='wb', suffix='.pdf', delete=False) as tmp:
            tmp.write(b"fake pdf content")
            tmp_path = tmp.name
        
        try:
            result = handler._extract_with_ocr_comprehensive(tmp_path, track_sources=True)
            # Должен обработать исключение и вернуть None
            assert result is None
        finally:
            if os.path.exists(tmp_path):
                os.remove(tmp_path)
    
    def test_load_pdf_aggregation_statistics(self):
        """Тест статистики агрегации данных"""
        pdf_path = self.test_dir / "examples" / "forwarder invoice sample.pdf"
        if pdf_path.exists():
            result = self.handler.load_pdf_document(str(pdf_path), track_sources=True)
            if result:
                df, sources_info = result
                if df is not None and not df.empty:
                    # Проверяем, что можно подсчитать статистику по источникам
                    if 'Source' in df.columns:
                        source_counts = df['Source'].value_counts()
                        assert len(source_counts) > 0
                        # Проверяем, что есть хотя бы один источник
                        assert source_counts.sum() == len(df)