tests/test_data_extractor.py - master - alsoalgo/CTF

tests/test_data_extractor.py
389 строк16 KB
alsoalgo
some improvements
15 дек 2025, 17:15
15 дек 2025, 17:155bd30a7
"""
Тесты для DataExtractor
"""
import pytest
import pandas as pd
import logging
import json
from pathlib import Path
from unittest.mock import Mock, patch, MagicMock

from src.components.data_extractor import DataExtractor


class TestDataExtractor:
    """Тесты для класса DataExtractor"""
    
    def setup_method(self):
        """Настройка перед каждым тестом"""
        self.extractor = DataExtractor()
        self.test_dir = Path(__file__).parent.parent
    
    def test_transform_dataframes_empty(self):
        """Тест преобразования пустого словаря"""
        result = self.extractor.transform_dataframes({})
        assert result == (None, None)
        
        result = self.extractor.transform_dataframes(None)
        assert result == (None, None)
    
    def test_transform_dataframes_single_document(self):
        """Тест преобразования одного документа"""
        # Создаем тестовый DataFrame (имитация invoice)
        test_df = pd.DataFrame({
            'Invoice Number': ['INV-001'],
            'Date': ['2024-01-15'],
            'Supplier': ['Test Supplier'],
            'Off-taker': ['Test Buyer'],
            'Product': ['Test Product'],
            'Quantity': ['100'],
            'Total Amount': ['10000.00'],
            'Currency': ['USD']
        })
        
        dataframes = {'test_invoice.csv': test_df}
        final_df, metadata_df = self.extractor.transform_dataframes(dataframes)
        
        assert final_df is not None
        assert metadata_df is not None
        assert isinstance(final_df, pd.DataFrame)
        assert isinstance(metadata_df, pd.DataFrame)
        assert 'Field' in final_df.columns
        assert 'PURCHASE' in final_df.columns
        assert 'SALE' in final_df.columns
        assert 'PURCHASE Confidence' in final_df.columns
        assert 'SALE Confidence' in final_df.columns
    
    def test_transform_dataframes_multiple_documents(self):
        """Тест преобразования нескольких документов"""
        # Invoice
        invoice_df = pd.DataFrame({
            'Invoice Number': ['INV-001'],
            'Supplier': ['Supplier A'],
            'Off-taker': ['Buyer B'],
            'Total Amount': ['5000.00'],
            'Currency': ['USD']
        })
        
        # Contract
        contract_df = pd.DataFrame({
            'Contract ID': ['CTF-001'],
            'Supplier': ['Supplier A'],
            'Off-taker': ['Buyer B'],
            'Total Value': ['5000.00'],
            'Currency': ['USD'],
            'Delivery Date': ['2024-03-15']
        })
        
        dataframes = {
            'invoice.csv': invoice_df,
            'contract.csv': contract_df
        }
        
        final_df, metadata_df = self.extractor.transform_dataframes(dataframes)
        
        assert final_df is not None
        assert metadata_df is not None
        
        # Проверяем, что данные из обоих документов попали в результат
        # Для Supplier используем PURCHASE, для Off-taker используем SALE
        supplier_row = final_df[final_df['Field'] == 'Supplier']
        if not supplier_row.empty:
            assert supplier_row.iloc[0]['PURCHASE'] != '' or supplier_row.iloc[0]['SALE'] != ''
    
    def test_detect_document_type(self):
        """Тест определения типа документа"""
        test_df = pd.DataFrame({'Column': ['Value']})
        
        # Тест по имени файла
        assert self.extractor._detect_document_type('invoice.pdf', test_df) == 'invoice'
        assert self.extractor._detect_document_type('contract.xlsx', test_df) == 'contract'
        assert self.extractor._detect_document_type('insurance_policy.pdf', test_df) == 'insurance'
        
        # Тест по колонкам
        invoice_df = pd.DataFrame({'Invoice Number': ['INV-001']})
        assert self.extractor._detect_document_type('unknown.csv', invoice_df) == 'invoice'
    
    def test_extract_field_value(self):
        """Тест извлечения значения поля"""
        test_df = pd.DataFrame({
            'Supplier': ['Test Supplier'],
            'Amount': ['1000.00']
        })
        
        value, row_idx, col_name, status, source_details = self.extractor._extract_field_value(
            test_df, 'supplier', 'invoice'
        )
        
        assert value == 'Test Supplier'
        assert row_idx == 0
        assert status in [0, 1, 2]
        # source_details может быть None или dict
        assert source_details is None or isinstance(source_details, dict)
    
    def test_create_template_dataframe(self):
        """Тест создания шаблона DataFrame"""
        template = self.extractor._create_template_dataframe()
        
        assert isinstance(template, pd.DataFrame)
        assert 'Field' in template.columns
        assert 'PURCHASE' in template.columns
        assert 'SALE' in template.columns
        assert 'PURCHASE Confidence' in template.columns
        assert 'SALE Confidence' in template.columns
        assert 'Supplier' in template['Field'].values
        assert 'Off-taker' in template['Field'].values
        
        # Проверяем, что все значения Confidence инициализированы нулями
        assert all(template['PURCHASE Confidence'] == 0)
        assert all(template['SALE Confidence'] == 0)
        assert len(template) == 15  # 15 полей в шаблоне (убрали FUND X, PURCHASE, SALE)
    
    def test_transform_dataframes_with_sources_info(self):
        """Тест преобразования с информацией об источниках (tuple)"""
        test_df = pd.DataFrame({
            'Supplier': ['Test Supplier'],
            'Off-taker': ['Test Buyer'],
            'Product': ['Test Product']
        })
        
        # Создаем sources_info
        sources_info = {
            'cells': {
                '0_Supplier': {
                    'page': 1,
                    'method': 'ocr',
                    'context': 'Supplier information'
                }
            }
        }
        
        dataframes = {'test_invoice.csv': (test_df, sources_info)}
        final_df, metadata_df = self.extractor.transform_dataframes(dataframes)
        
        assert final_df is not None
        assert metadata_df is not None
        # Проверяем, что метаданные содержат информацию об источниках
        assert 'source_page' in metadata_df.columns
        assert 'source_context' in metadata_df.columns
        assert 'extraction_method' in metadata_df.columns
    
    def test_transform_dataframes_backward_compatibility(self):
        """Тест обратной совместимости (DataFrame без tuple)"""
        test_df = pd.DataFrame({
            'Supplier': ['Test Supplier'],
            'Off-taker': ['Test Buyer']
        })
        
        # Передаем просто DataFrame, без tuple
        dataframes = {'test_invoice.csv': test_df}
        final_df, metadata_df = self.extractor.transform_dataframes(dataframes)
        
        assert final_df is not None
        assert metadata_df is not None
    
    def test_extract_field_value_with_sources_info(self):
        """Тест извлечения значения поля с sources_info"""
        test_df = pd.DataFrame({
            'Supplier': ['Test Supplier']
        })
        
        sources_info = {
            'cells': {
                '0_Supplier': {
                    'page': 1,
                    'method': 'ocr',
                    'context': 'Found in header'
                }
            }
        }
        
        value, row_idx, col_name, status, source_details = self.extractor._extract_field_value(
            test_df, 'supplier', 'invoice', sources_info=sources_info
        )
        
        assert value == 'Test Supplier'
        assert source_details is not None
        assert isinstance(source_details, dict)
        assert source_details.get('page') == 1
        assert source_details.get('method') == 'ocr'
    
    def test_extract_field_value_without_sources_info(self):
        """Тест извлечения значения поля без sources_info"""
        test_df = pd.DataFrame({
            'Supplier': ['Test Supplier']
        })
        
        value, row_idx, col_name, status, source_details = self.extractor._extract_field_value(
            test_df, 'supplier', 'invoice', sources_info=None
        )
        
        assert value == 'Test Supplier'
        assert source_details is None
    
    def test_transform_dataframes_logging(self, caplog):
        """Тест логирования при преобразовании данных"""
        logging.basicConfig(level=logging.INFO)
        
        test_df = pd.DataFrame({
            'Supplier': ['Test Supplier'],
            'Off-taker': ['Test Buyer'],
            'Product': ['Test Product']
        })
        
        dataframes = {'test_invoice.csv': test_df}
        
        with caplog.at_level(logging.INFO):
            final_df, metadata_df = self.extractor.transform_dataframes(dataframes)
            # Проверяем, что есть логи о обработке документов
            log_messages = [record.message for record in caplog.records]
            assert any("Обрабатываю документ" in msg for msg in log_messages)
    
    def test_extract_field_value_logging(self, caplog):
        """Тест логирования при извлечении полей"""
        logging.basicConfig(level=logging.INFO)
        
        test_df = pd.DataFrame({
            'Supplier': ['Test Supplier']
        })
        
        sources_info = {
            'cells': {
                '0_Supplier': {
                    'page': 1,
                    'method': 'ocr',
                    'context': 'Found in header'
                }
            }
        }
        
        with caplog.at_level(logging.INFO):
            value, row_idx, col_name, status, source_details = self.extractor._extract_field_value(
                test_df, 'supplier', 'invoice', sources_info=sources_info
            )
            # Логи могут быть или не быть, проверяем что нет ошибок
            assert value == 'Test Supplier'
    
    def test_parse_llm_response_with_confidence(self):
        """Тест парсинга ответа LLM с confidence"""
        # Создаем мок ответа LLM с новой структурой (purchase_value, sale_value, purchase_confidence, sale_confidence)
        llm_response = json.dumps({
            "fields": {
                "Supplier": {
                    "purchase_value": "Test Supplier",
                    "sale_value": "",
                    "source_file": "test.pdf",
                    "source_method": "ocr",
                    "purchase_confidence": 95,
                    "sale_confidence": 0
                },
                "Off-taker": {
                    "purchase_value": "",
                    "sale_value": "Test Buyer",
                    "source_file": "test.pdf",
                    "source_method": "table",
                    "purchase_confidence": 0,
                    "sale_confidence": 80
                },
                "Goods": {
                    "purchase_value": "",
                    "sale_value": "",
                    "source_file": "",
                    "source_method": "",
                    "purchase_confidence": 0,
                    "sale_confidence": 0
                }
            }
        })
        
        template = self.extractor._create_template_dataframe()
        metadata = pd.DataFrame()
        dataframes = {}
        
        # Вызываем _parse_llm_response напрямую
        result_template, result_metadata = self.extractor._parse_llm_response(
            llm_response, template, metadata, dataframes
        )
        
        # Проверяем, что confidence сохранен
        supplier_row = result_template[result_template['Field'] == 'Supplier']
        assert len(supplier_row) > 0
        # Supplier должен быть в PURCHASE
        assert supplier_row.iloc[0]['PURCHASE Confidence'] == 95
        
        off_taker_row = result_template[result_template['Field'] == 'Off-taker']
        assert len(off_taker_row) > 0
        # Off-taker должен быть в SALE
        assert off_taker_row.iloc[0]['SALE Confidence'] == 80
        
        goods_row = result_template[result_template['Field'] == 'Goods']
        if len(goods_row) > 0:
            # Goods может быть в PURCHASE или SALE, проверяем оба
            assert goods_row.iloc[0]['PURCHASE Confidence'] == 0 or goods_row.iloc[0]['SALE Confidence'] == 0
    
    def test_parse_llm_response_confidence_status_mapping(self):
        """Тест маппинга confidence в status для метаданных"""
        llm_response = json.dumps({
            "fields": {
                "Supplier": {
                    "purchase_value": "Test Supplier",
                    "sale_value": "",
                    "source_file": "test.pdf",
                    "source_method": "ocr",
                    "purchase_confidence": 95,  # Высокая уверенность -> status 0
                    "sale_confidence": 0
                },
                "Goods": {
                    "purchase_value": "Test Product",
                    "sale_value": "",
                    "source_file": "test.pdf",
                    "source_method": "ocr",
                    "purchase_confidence": 50,  # Средняя уверенность -> status 1
                    "sale_confidence": 0
                },
                "Amount, USD": {
                    "purchase_value": "1000",
                    "sale_value": "",
                    "source_file": "test.pdf",
                    "source_method": "ocr",
                    "purchase_confidence": 20,  # Низкая уверенность -> status 2
                    "sale_confidence": 0
                }
            }
        })
        
        template = self.extractor._create_template_dataframe()
        metadata = pd.DataFrame()
        dataframes = {}
        
        result_template, result_metadata = self.extractor._parse_llm_response(
            llm_response, template, metadata, dataframes
        )
        
        # Проверяем, что status правильно установлен на основе confidence
        # status 0 = высокая уверенность (>=70)
        # status 1 = средняя уверенность (30-69)
        # status 2 = низкая уверенность (<30)
        
        supplier_idx = result_template[result_template['Field'] == 'Supplier'].index[0]
        # Supplier должен быть в PURCHASE
        supplier_meta = result_metadata.loc[(supplier_idx, 'PURCHASE')]
        assert supplier_meta['status'] == 0  # 95 >= 70
        
        goods_idx = result_template[result_template['Field'] == 'Goods'].index[0]
        # Goods может быть в PURCHASE или SALE, проверяем оба
        if (goods_idx, 'PURCHASE') in result_metadata.index:
            goods_meta = result_metadata.loc[(goods_idx, 'PURCHASE')]
            assert goods_meta['status'] == 1  # 50 в диапазоне 30-69
        elif (goods_idx, 'SALE') in result_metadata.index:
            goods_meta = result_metadata.loc[(goods_idx, 'SALE')]
            assert goods_meta['status'] == 1  # 50 в диапазоне 30-69
        
        amount_idx = result_template[result_template['Field'] == 'Amount, USD'].index[0]
        # Amount может быть в PURCHASE или SALE, проверяем оба
        if (amount_idx, 'PURCHASE') in result_metadata.index:
            amount_meta = result_metadata.loc[(amount_idx, 'PURCHASE')]
            assert amount_meta['status'] == 2  # 20 < 30
        elif (amount_idx, 'SALE') in result_metadata.index:
            amount_meta = result_metadata.loc[(amount_idx, 'SALE')]
            assert amount_meta['status'] == 2  # 20 < 30