""" Unit tests for ParseDocumentOperator. """ import pytest from unittest.mock import Mock, MagicMock, patch from plugins.operators.parse_document import ParseDocumentOperator class TestParseDocumentOperator: """Test suite for ParseDocumentOperator.""" @pytest.fixture def mock_context(self): """Create a mock Airflow context.""" context = { 'ti': Mock(), } return context @pytest.fixture def sample_fetch_result(self): """Sample fetch result data.""" return { 'local_path': '/tmp/test/sample.xlsx', 'content_type': 'application/vnd.openxmlformats-officedocument.spreadsheetml.sheet', 'original_filename': 'sample.xlsx', 'file_extension': 'xlsx', 'document_id': 'test-doc-123', } def test_parse_excel_document_normal_tables(self, mock_context, sample_fetch_result): """Test parsing Excel document with normal tables (no large tables).""" mock_context['ti'].xcom_pull.return_value = sample_fetch_result with patch('plugins.operators.parse_document.DocumentToHtmlConverter') as mock_converter_class, \ patch('plugins.operators.parse_document.HtmlTableExtractor') as mock_extractor_class, \ patch('plugins.operators.parse_document.LargeTableExtractor') as mock_large_extractor_class: # Mock large table extractor - returns normal tables only mock_large_extractor = MagicMock() mock_large_extractor.detect_tables.return_value = [ {'sheet': 'Sheet1', 'is_large': False, 'row_count': 50} ] mock_large_extractor_class.return_value = mock_large_extractor # Mock converter mock_converter = MagicMock() mock_converter.convert_to_html.return_value = '...
' mock_converter_class.return_value = mock_converter # Mock table extractor mock_table1 = MagicMock() mock_table1.html = 'Table 1
' mock_table1.rows = [['A', 'B'], ['1', '2']] mock_raw_table = MagicMock() mock_raw_table.rows = [['Raw', 'Data']] mock_extractor = MagicMock() mock_extractor.extract_all_tables.return_value = ([mock_table1], [mock_raw_table]) mock_extractor_class.return_value = mock_extractor op = ParseDocumentOperator( task_id='test_parse', doc_type='excel', fetch_document_task_id='fetch_document' ) result = op.execute(mock_context) # Verify XCom pull mock_context['ti'].xcom_pull.assert_called_once_with(task_ids='fetch_document') # Verify large table detection was called mock_large_extractor.detect_tables.assert_called_once_with('/tmp/test/sample.xlsx') # Verify converter was used (for normal tables) mock_converter.convert_to_html.assert_called_once_with('/tmp/test/sample.xlsx') # Verify extractor was used mock_extractor.extract_all_tables.assert_called_once() # Verify result structure assert result['doc_type'] == 'excel' assert result['html_content'] == '...
' assert result['markdown_content'] is None assert len(result['tables']) == 1 assert result['tables'][0]['html'] == 'Table 1
' assert len(result['raw_text_tables']) == 1 assert result['large_table_records'] == [] assert result['document_id'] == 'test-doc-123' assert result['local_path'] == '/tmp/test/sample.xlsx' def test_parse_excel_document_with_large_tables(self, mock_context, sample_fetch_result): """Test parsing Excel document with large tables (>1000 rows).""" mock_context['ti'].xcom_pull.return_value = sample_fetch_result with patch('plugins.operators.parse_document.DocumentToHtmlConverter') as mock_converter_class, \ patch('plugins.operators.parse_document.HtmlTableExtractor') as mock_extractor_class, \ patch('plugins.operators.parse_document.LargeTableExtractor') as mock_large_extractor_class: # Mock large table extractor - returns large tables only mock_large_result = MagicMock() mock_large_result.success = True mock_large_result.records = [ {'row': 1, 'data': 'large_record_1'}, {'row': 2, 'data': 'large_record_2'}, ] mock_large_result.rows_processed = 1500 mock_large_result.sheet_name = 'LargeSheet' mock_large_extractor = MagicMock() mock_large_extractor.detect_tables.return_value = [ {'sheet': 'LargeSheet', 'is_large': True, 'row_count': 1500} ] mock_large_extractor.extract_table.return_value = mock_large_result mock_large_extractor_class.return_value = mock_large_extractor # Converter and extractor should NOT be called (no normal tables) mock_converter = MagicMock() mock_converter_class.return_value = mock_converter mock_extractor = MagicMock() mock_extractor_class.return_value = mock_extractor op = ParseDocumentOperator( task_id='test_parse', doc_type='excel', ) result = op.execute(mock_context) # Verify large table detection was called mock_large_extractor.detect_tables.assert_called_once_with('/tmp/test/sample.xlsx') # Verify large table extraction was called mock_large_extractor.extract_table.assert_called_once() # Verify converter was NOT called (only large tables) mock_converter.convert_to_html.assert_not_called() # Verify result structure assert result['doc_type'] == 'excel' assert result['html_content'] is None # No Docling conversion for large-only assert result['tables'] == [] assert result['raw_text_tables'] == [] assert len(result['large_table_records']) == 2 assert result['large_table_records'][0]['data'] == 'large_record_1' def test_parse_excel_document_mixed_tables(self, mock_context, sample_fetch_result): """Test parsing Excel document with both large and normal tables.""" mock_context['ti'].xcom_pull.return_value = sample_fetch_result with patch('plugins.operators.parse_document.DocumentToHtmlConverter') as mock_converter_class, \ patch('plugins.operators.parse_document.HtmlTableExtractor') as mock_extractor_class, \ patch('plugins.operators.parse_document.LargeTableExtractor') as mock_large_extractor_class: # Mock large table extractor - returns both large and normal mock_large_result = MagicMock() mock_large_result.success = True mock_large_result.records = [{'data': 'large_record'}] mock_large_result.rows_processed = 2000 mock_large_result.sheet_name = 'LargeSheet' mock_large_extractor = MagicMock() mock_large_extractor.detect_tables.return_value = [ {'sheet': 'LargeSheet', 'is_large': True, 'row_count': 2000}, {'sheet': 'NormalSheet', 'is_large': False, 'row_count': 100}, ] mock_large_extractor.extract_table.return_value = mock_large_result mock_large_extractor_class.return_value = mock_large_extractor # Mock converter for normal tables mock_converter = MagicMock() mock_converter.convert_to_html.return_value = 'normal
' mock_converter_class.return_value = mock_converter # Mock table extractor mock_table = MagicMock() mock_table.html = 'Normal Table
' mock_table.rows = [['A'], ['1']] mock_extractor = MagicMock() mock_extractor.extract_all_tables.return_value = ([mock_table], []) mock_extractor_class.return_value = mock_extractor op = ParseDocumentOperator(task_id='test_parse', doc_type='excel') result = op.execute(mock_context) # Verify both paths were executed mock_large_extractor.detect_tables.assert_called_once() mock_large_extractor.extract_table.assert_called_once() mock_converter.convert_to_html.assert_called_once() mock_extractor.extract_all_tables.assert_called_once() # Verify result includes both assert len(result['large_table_records']) == 1 assert len(result['tables']) == 1 assert result['html_content'] == 'normal
' def test_parse_excel_large_table_extraction_failure(self, mock_context, sample_fetch_result): """Test handling of large table extraction failure.""" mock_context['ti'].xcom_pull.return_value = sample_fetch_result with patch('plugins.operators.parse_document.DocumentToHtmlConverter') as mock_converter_class, \ patch('plugins.operators.parse_document.HtmlTableExtractor') as mock_extractor_class, \ patch('plugins.operators.parse_document.LargeTableExtractor') as mock_large_extractor_class: # Mock large table extractor - extraction fails mock_large_result = MagicMock() mock_large_result.success = False mock_large_result.error = 'Extraction failed' mock_large_extractor = MagicMock() mock_large_extractor.detect_tables.return_value = [ {'sheet': 'LargeSheet', 'is_large': True, 'row_count': 1500} ] mock_large_extractor.extract_table.return_value = mock_large_result mock_large_extractor_class.return_value = mock_large_extractor mock_converter = MagicMock() mock_converter_class.return_value = mock_converter mock_extractor = MagicMock() mock_extractor_class.return_value = mock_extractor op = ParseDocumentOperator(task_id='test_parse', doc_type='excel') result = op.execute(mock_context) # Verify extraction was attempted mock_large_extractor.extract_table.assert_called_once() # Result should have empty large_table_records due to failure assert result['large_table_records'] == [] def test_parse_excel_no_tables_detected(self, mock_context, sample_fetch_result): """Test Excel parsing when border detection finds no tables.""" mock_context['ti'].xcom_pull.return_value = sample_fetch_result with patch('plugins.operators.parse_document.DocumentToHtmlConverter') as mock_converter_class, \ patch('plugins.operators.parse_document.HtmlTableExtractor') as mock_extractor_class, \ patch('plugins.operators.parse_document.LargeTableExtractor') as mock_large_extractor_class: # Mock large table extractor - no tables detected mock_large_extractor = MagicMock() mock_large_extractor.detect_tables.return_value = [] mock_large_extractor_class.return_value = mock_large_extractor # Should fall back to Docling conversion mock_converter = MagicMock() mock_converter.convert_to_html.return_value = 'fallback' mock_converter_class.return_value = mock_converter mock_extractor = MagicMock() mock_extractor.extract_all_tables.return_value = ([], []) mock_extractor_class.return_value = mock_extractor op = ParseDocumentOperator(task_id='test_parse', doc_type='excel') result = op.execute(mock_context) # Verify fallback to Docling was used mock_converter.convert_to_html.assert_called_once() assert result['html_content'] == 'fallback' def test_parse_word_document(self, mock_context, sample_fetch_result): """Test parsing Word document to Markdown.""" fetch_result = sample_fetch_result.copy() fetch_result['file_extension'] = 'docx' fetch_result['local_path'] = '/tmp/test/document.docx' mock_context['ti'].xcom_pull.return_value = fetch_result with patch('plugins.operators.parse_document.DocumentToHtmlConverter') as mock_converter_class: mock_converter = MagicMock() mock_converter.convert_to_html.return_value = '# Document Title\n\nContent here...' mock_converter_class.return_value = mock_converter op = ParseDocumentOperator( task_id='test_parse', doc_type='word' ) result = op.execute(mock_context) # Verify converter was used mock_converter.convert_to_html.assert_called_once_with('/tmp/test/document.docx') # Verify result structure assert result['doc_type'] == 'word' assert result['html_content'] is None assert result['markdown_content'] == '# Document Title\n\nContent here...' assert result['tables'] == [] assert result['raw_text_tables'] == [] assert result['document_id'] == 'test-doc-123' def test_parse_pdf_document(self, mock_context, sample_fetch_result): """Test parsing PDF document to Markdown.""" fetch_result = sample_fetch_result.copy() fetch_result['file_extension'] = 'pdf' fetch_result['local_path'] = '/tmp/test/document.pdf' mock_context['ti'].xcom_pull.return_value = fetch_result with patch('plugins.operators.parse_document.DocumentToHtmlConverter') as mock_converter_class: mock_converter = MagicMock() mock_converter.convert_to_html.return_value = '# PDF Content\n\nParsed text' mock_converter_class.return_value = mock_converter op = ParseDocumentOperator( task_id='test_parse', doc_type='pdf' ) result = op.execute(mock_context) assert result['doc_type'] == 'pdf' assert result['markdown_content'] == '# PDF Content\n\nParsed text' def test_parse_powerpoint_document(self, mock_context, sample_fetch_result): """Test parsing PowerPoint document to Markdown.""" fetch_result = sample_fetch_result.copy() fetch_result['file_extension'] = 'pptx' fetch_result['local_path'] = '/tmp/test/presentation.pptx' mock_context['ti'].xcom_pull.return_value = fetch_result with patch('plugins.operators.parse_document.DocumentToHtmlConverter') as mock_converter_class: mock_converter = MagicMock() mock_converter.convert_to_html.return_value = '# Slide 1\n\nContent' mock_converter_class.return_value = mock_converter op = ParseDocumentOperator( task_id='test_parse', doc_type='powerpoint' ) result = op.execute(mock_context) assert result['doc_type'] == 'powerpoint' assert result['markdown_content'] == '# Slide 1\n\nContent' def test_parse_no_fetch_result_raises(self, mock_context): """Test that missing fetch result raises ValueError.""" mock_context['ti'].xcom_pull.return_value = None op = ParseDocumentOperator(task_id='test_parse', doc_type='excel') with pytest.raises(ValueError, match='No fetch result found'): op.execute(mock_context) def test_parse_conversion_failure_raises(self, mock_context, sample_fetch_result): """Test that conversion failure raises exception.""" mock_context['ti'].xcom_pull.return_value = sample_fetch_result with patch('plugins.operators.parse_document.DocumentToHtmlConverter') as mock_converter_class, \ patch('plugins.operators.parse_document.LargeTableExtractor') as mock_large_extractor_class: # Mock large table extractor - returns normal table to trigger Docling mock_large_extractor = MagicMock() mock_large_extractor.detect_tables.return_value = [ {'sheet': 'Sheet1', 'is_large': False} ] mock_large_extractor_class.return_value = mock_large_extractor mock_converter = MagicMock() mock_converter.convert_to_html.side_effect = Exception('Conversion failed') mock_converter_class.return_value = mock_converter op = ParseDocumentOperator(task_id='test_parse', doc_type='excel') with pytest.raises(Exception, match='Conversion failed'): op.execute(mock_context) def test_custom_fetch_task_id(self, mock_context, sample_fetch_result): """Test using custom fetch task ID.""" mock_context['ti'].xcom_pull.return_value = sample_fetch_result with patch('plugins.operators.parse_document.DocumentToHtmlConverter') as mock_converter_class: mock_converter = MagicMock() mock_converter.convert_to_html.return_value = 'Content' mock_converter_class.return_value = mock_converter op = ParseDocumentOperator( task_id='test_parse', doc_type='word', fetch_document_task_id='custom_fetch' ) result = op.execute(mock_context) # Verify custom task_id was used mock_context['ti'].xcom_pull.assert_called_once_with(task_ids='custom_fetch') assert result['doc_type'] == 'word' def test_excel_with_no_tables(self, mock_context, sample_fetch_result): """Test Excel parsing when no tables are found (after Docling conversion).""" mock_context['ti'].xcom_pull.return_value = sample_fetch_result with patch('plugins.operators.parse_document.DocumentToHtmlConverter') as mock_converter_class, \ patch('plugins.operators.parse_document.HtmlTableExtractor') as mock_extractor_class, \ patch('plugins.operators.parse_document.LargeTableExtractor') as mock_large_extractor_class: # Large table extractor finds normal tables mock_large_extractor = MagicMock() mock_large_extractor.detect_tables.return_value = [ {'sheet': 'Sheet1', 'is_large': False} ] mock_large_extractor_class.return_value = mock_large_extractor mock_converter = MagicMock() mock_converter.convert_to_html.return_value = 'No tables' mock_converter_class.return_value = mock_converter mock_extractor = MagicMock() mock_extractor.extract_all_tables.return_value = ([], []) # No tables mock_extractor_class.return_value = mock_extractor op = ParseDocumentOperator(task_id='test_parse', doc_type='excel') result = op.execute(mock_context) assert result['tables'] == [] assert result['raw_text_tables'] == [] assert result['html_content'] == 'No tables' assert result['large_table_records'] == []