""" Unit tests for Word Document Preprocessing. Tests cover: - Image extraction from Word documents - AI image captioning via GPT-4o vision - Shape/text box extraction - Comment extraction - Inline embedding of all content """ import pytest from pathlib import Path from unittest.mock import Mock, MagicMock, patch from docx import Document from src.extraction_v2.word_preprocessor import WordPreprocessor class TestWordPreprocessor: """Test suite for WordPreprocessor class.""" @pytest.fixture def mock_openai_client(self): """Create a mock Azure OpenAI client.""" client = Mock() response = Mock() response.choices = [Mock(message=Mock(content="Test image caption"))] client.chat.completions.create.return_value = response return client @pytest.fixture def preprocessor(self, mock_openai_client): """Create a WordPreprocessor instance with mocked client.""" return WordPreprocessor(openai_client=mock_openai_client) def test_extract_images_empty_document(self, preprocessor): """Test image extraction on document with no images.""" # Create empty document doc = Document() doc.add_paragraph("This is a test paragraph with no images.") images = preprocessor._extract_images(doc) assert images == [] def test_get_paragraph_context(self, preprocessor): """Test paragraph context extraction.""" doc = Document() doc.add_paragraph("Paragraph 1") doc.add_paragraph("Paragraph 2") doc.add_paragraph("Paragraph 3 - Target") doc.add_paragraph("Paragraph 4") doc.add_paragraph("Paragraph 5") # Get context around paragraph 2 (index 2) with window=2 context = preprocessor._get_paragraph_context(doc, para_idx=2, window=2) # Should include paragraphs 0, 1, 2, 3, 4 assert "Paragraph 1" in context assert "Paragraph 2" in context assert "Paragraph 3 - Target" in context assert "Paragraph 4" in context assert "Paragraph 5" in context def test_get_paragraph_context_edge_case_start(self, preprocessor): """Test paragraph context at document start.""" doc = Document() doc.add_paragraph("First paragraph") doc.add_paragraph("Second paragraph") doc.add_paragraph("Third paragraph") # Get context at start (index 0) context = preprocessor._get_paragraph_context(doc, para_idx=0, window=2) assert "First paragraph" in context assert "Second paragraph" in context assert "Third paragraph" in context def test_get_paragraph_context_edge_case_end(self, preprocessor): """Test paragraph context at document end.""" doc = Document() doc.add_paragraph("First paragraph") doc.add_paragraph("Second paragraph") doc.add_paragraph("Last paragraph") # Get context at end (index 2) context = preprocessor._get_paragraph_context(doc, para_idx=2, window=2) assert "First paragraph" in context assert "Second paragraph" in context assert "Last paragraph" in context def test_detect_image_type_png(self, preprocessor): """Test PNG image type detection.""" png_data = b'\x89PNG\r\n\x1a\n' assert preprocessor._detect_image_type(png_data) == "png" def test_detect_image_type_jpeg(self, preprocessor): """Test JPEG image type detection.""" jpeg_data = b'\xFF\xD8\xFF\xE0' assert preprocessor._detect_image_type(jpeg_data) == "jpeg" def test_detect_image_type_gif(self, preprocessor): """Test GIF image type detection.""" gif_data = b'GIF89a' assert preprocessor._detect_image_type(gif_data) == "gif" def test_detect_image_type_bmp(self, preprocessor): """Test BMP image type detection.""" bmp_data = b'BM' assert preprocessor._detect_image_type(bmp_data) == "bmp" def test_detect_image_type_unknown(self, preprocessor): """Test unknown image type defaults to PNG.""" unknown_data = b'UNKNOWN' assert preprocessor._detect_image_type(unknown_data) == "png" def test_generate_image_captions_with_mock(self, preprocessor, mock_openai_client): """Test AI image caption generation with mock client.""" images = [ { 'data': b'\x89PNG test data', 'context': 'Test context around image', 'para_idx': 0 } ] captions = preprocessor._generate_image_captions(images) assert len(captions) == 1 assert captions[0] == "Test image caption" assert mock_openai_client.chat.completions.create.called def test_generate_image_captions_error_handling(self, preprocessor): """Test caption generation error handling.""" # Create preprocessor with client that raises error error_client = Mock() error_client.chat.completions.create.side_effect = Exception("API Error") preprocessor._client = error_client images = [ { 'data': b'\x89PNG test data', 'context': 'Test context', 'para_idx': 0 } ] captions = preprocessor._generate_image_captions(images) assert len(captions) == 1 assert "[Image: Caption unavailable]" in captions[0] def test_embed_image_captions(self, preprocessor): """Test inline embedding of image captions.""" doc = Document() para = doc.add_paragraph("Text before image. ") images = [ { 'paragraph': para, 'para_idx': 0 } ] captions = ["A beautiful sunset over the ocean"] preprocessor._embed_image_captions(doc, images, captions) # Check that caption was added to paragraph para_text = para.text assert "[IMAGE: A beautiful sunset over the ocean]" in para_text def test_extract_shapes_empty_document(self, preprocessor): """Test shape extraction on document with no shapes.""" doc = Document() doc.add_paragraph("Simple text, no shapes.") shapes = preprocessor._extract_shapes(doc) assert shapes == [] def test_extract_textbox_text_no_textbox(self, preprocessor): """Test textbox text extraction on element without textbox.""" doc = Document() para = doc.add_paragraph("Regular paragraph") # Element without textbox content text = preprocessor._extract_textbox_text(para._element) # Should return None or empty string assert text is None or text == "" def test_extract_comments_empty_document(self, preprocessor): """Test comment extraction on document with no comments.""" doc = Document() doc.add_paragraph("Text without comments") comments = preprocessor._extract_comments(doc) assert comments == [] def test_apply_comment_markers_empty_list(self, preprocessor): """Test applying comment markers with empty list.""" doc = Document() initial_para_count = len(doc.paragraphs) preprocessor._apply_comment_markers(doc, []) # Should not add anything if no comments assert len(doc.paragraphs) == initial_para_count def test_apply_comment_markers_with_comments(self, preprocessor): """Test applying comment markers with actual comments.""" doc = Document() initial_para_count = len(doc.paragraphs) comments = [ { 'id': '1', 'author': 'John Doe', 'text': 'This is a test comment', 'date': '2025-12-01' }, { 'id': '2', 'author': 'Jane Smith', 'text': 'Another comment here', 'date': '2025-12-01' } ] preprocessor._apply_comment_markers(doc, comments) # Should add: # 1 blank paragraph + 1 "DOCUMENT COMMENTS:" + 2 comment paragraphs = 4 new paragraphs assert len(doc.paragraphs) == initial_para_count + 4 # Check comment markers were added doc_text = "\n".join([p.text for p in doc.paragraphs]) assert "DOCUMENT COMMENTS:" in doc_text assert "[COMMENT_1]" in doc_text assert "John Doe" in doc_text assert "This is a test comment" in doc_text assert "[COMMENT_2]" in doc_text assert "Jane Smith" in doc_text assert "Another comment here" in doc_text def test_embed_shape_descriptions(self, preprocessor): """Test inline embedding of shape descriptions.""" doc = Document() para = doc.add_paragraph("Text near shape. ") shapes = [ { 'type': 'textbox', 'text': 'Important note in text box', 'para_idx': 0, 'paragraph': para } ] preprocessor._embed_shape_descriptions(doc, shapes) # Check that shape description was added to paragraph para_text = para.text assert "[SHAPE: Important note in text box]" in para_text @patch('src.extraction_v2.word_preprocessor.Document') def test_preprocess_full_workflow(self, mock_document_class, preprocessor, tmp_path): """Test full preprocessing workflow.""" # Create a test file test_file = tmp_path / "test_doc.docx" test_file.touch() # Mock Document mock_doc = MagicMock() mock_doc.paragraphs = [] mock_document_class.return_value = mock_doc # Mock the extraction methods to return empty lists preprocessor._extract_images = Mock(return_value=[]) preprocessor._extract_shapes = Mock(return_value=[]) preprocessor._extract_comments = Mock(return_value=[]) # Run preprocessing output_path = preprocessor.preprocess(test_file) # Verify preprocessing steps were called assert preprocessor._extract_images.called assert preprocessor._extract_shapes.called assert preprocessor._extract_comments.called # Verify output path assert output_path is not None assert "_preprocessed" in str(output_path) def test_lazy_client_initialization(self): """Test that OpenAI client is initialized lazily.""" preprocessor = WordPreprocessor(openai_client=None) assert preprocessor._client is None # Client should be created on first use with patch('openai.AzureOpenAI') as mock_azure: with patch('src.core.config.settings') as mock_settings: mock_settings.azure_openai_api_key = "test_key" mock_settings.azure_openai_api_version = "2024-02-15-preview" mock_settings.azure_openai_endpoint = "https://test.openai.azure.com/" client = preprocessor._get_client() assert client is not None assert mock_azure.called class TestWordPreprocessorIntegration: """Integration tests for WordPreprocessor with real documents.""" def test_preprocess_with_real_empty_document(self, tmp_path): """Test preprocessing on a real but empty document.""" # Create a simple document doc = Document() doc.add_paragraph("This is a simple test document.") doc.add_paragraph("It has no images, shapes, or comments.") # Save it test_file = tmp_path / "empty_test.docx" doc.save(str(test_file)) # Preprocess it mock_client = Mock() preprocessor = WordPreprocessor(openai_client=mock_client) output_path = preprocessor.preprocess(test_file) # Verify output exists assert output_path.exists() assert output_path.name.endswith("_preprocessed.docx") # Verify the preprocessed document can be opened preprocessed_doc = Document(str(output_path)) assert len(preprocessed_doc.paragraphs) >= 2 if __name__ == "__main__": pytest.main([__file__, "-v"])