"""
Integration tests for V2 extraction pipeline with raw text handling.
Tests the full pipeline behavior when processing mixed content
(tables and raw text).
"""
import pytest
from unittest.mock import Mock, patch, MagicMock, AsyncMock
from uuid import uuid4
from src.extraction_v2.pipeline import ExtractionPipelineV2, ExtractionResult
from src.extraction_v2.html_table_extractor import HtmlTable, HtmlTableExtractor
from src.extraction_v2.record_builder import RecordBuilder
from src.extraction_v2.llm_header_detector import HeaderStructure
class TestPipelineRawTextIntegration:
"""Integration tests for pipeline raw text handling."""
@pytest.fixture
def mock_converter(self):
"""Create a mock converter."""
converter = Mock()
converter.convert_to_html.return_value = ""
converter.get_sheet_name_for_table.return_value = "TestSheet"
return converter
@pytest.fixture
def mock_detector(self):
"""Create a mock LLM header detector."""
detector = Mock()
detector.detect_headers.return_value = HeaderStructure(
header_rows=[0],
headers=[["Col1", "Col2"]],
)
return detector
@pytest.mark.asyncio
async def test_pipeline_processes_mixed_content(
self, mock_converter, mock_detector
):
"""Test pipeline correctly processes both valid tables and raw text."""
# Create mixed HTML content
mixed_html = """
| Raw text content line 1 |
| Raw text content line 2 |
"""
mock_converter.convert_to_html.return_value = mixed_html
pipeline = ExtractionPipelineV2(
converter=mock_converter,
detector=mock_detector,
)
document_id = uuid4()
result = await pipeline.process_document(document_id, "test.xlsx")
assert result.success is True
# Should have records from both tables
assert result.record_count > 0
# Table count should include both types
assert result.table_count == 2
# Check that we have both types of records
raw_text_records = [
r for r in result.records if r.get("_source", {}).get("row") == -1
]
table_records = [
r for r in result.records if r.get("_source", {}).get("row") != -1
]
assert len(raw_text_records) >= 1 # At least one raw text record
assert raw_text_records[0]["content"]["text"] == "Raw text content line 1\nRaw text content line 2"
@pytest.mark.asyncio
async def test_pipeline_handles_only_raw_text(self, mock_converter, mock_detector):
"""Test pipeline when document contains only raw text tables."""
raw_text_only_html = """
| First paragraph text |
| Second paragraph text |
"""
mock_converter.convert_to_html.return_value = raw_text_only_html
pipeline = ExtractionPipelineV2(
converter=mock_converter,
detector=mock_detector,
)
document_id = uuid4()
result = await pipeline.process_document(document_id, "test.xlsx")
assert result.success is True
assert result.table_count == 2
# All records should be raw text (row = -1)
for record in result.records:
assert record["_source"]["row"] == -1
# LLM detector should NOT have been called (raw text skips LLM)
mock_detector.detect_headers.assert_not_called()
@pytest.mark.asyncio
async def test_pipeline_handles_only_tables(self, mock_converter, mock_detector):
"""Test pipeline when document contains only valid data tables."""
tables_only_html = """
"""
mock_converter.convert_to_html.return_value = tables_only_html
pipeline = ExtractionPipelineV2(
converter=mock_converter,
detector=mock_detector,
)
document_id = uuid4()
result = await pipeline.process_document(document_id, "test.xlsx")
assert result.success is True
assert result.table_count == 2
# No records should have row = -1 (no raw text)
for record in result.records:
assert record.get("_source", {}).get("row", 0) != -1
# LLM detector should have been called for each table
assert mock_detector.detect_headers.call_count == 2
@pytest.mark.asyncio
async def test_pipeline_no_llm_for_raw_text(self, mock_converter, mock_detector):
"""Test that LLM is NOT called for raw text tables (cost savings)."""
# One raw text table
html = """
"""
mock_converter.convert_to_html.return_value = html
pipeline = ExtractionPipelineV2(
converter=mock_converter,
detector=mock_detector,
)
document_id = uuid4()
result = await pipeline.process_document(document_id, "test.xlsx")
assert result.success is True
# LLM should NOT be called for raw text
mock_detector.detect_headers.assert_not_called()
@pytest.mark.asyncio
async def test_pipeline_llm_only_for_valid_tables(
self, mock_converter, mock_detector
):
"""Test that LLM is only called for valid data tables."""
html = """
"""
mock_converter.convert_to_html.return_value = html
pipeline = ExtractionPipelineV2(
converter=mock_converter,
detector=mock_detector,
)
document_id = uuid4()
result = await pipeline.process_document(document_id, "test.xlsx")
assert result.success is True
# LLM should only be called for the 2 valid tables (not 4 total)
assert mock_detector.detect_headers.call_count == 2
@pytest.mark.asyncio
async def test_pipeline_empty_document(self, mock_converter, mock_detector):
"""Test pipeline with no tables at all."""
mock_converter.convert_to_html.return_value = ""
pipeline = ExtractionPipelineV2(
converter=mock_converter,
detector=mock_detector,
)
document_id = uuid4()
result = await pipeline.process_document(document_id, "test.xlsx")
assert result.success is True
assert result.record_count == 0
assert result.table_count == 0
@pytest.mark.asyncio
async def test_raw_text_record_structure(self, mock_converter, mock_detector):
"""Test the structure of raw text records."""
html = """
"""
mock_converter.convert_to_html.return_value = html
pipeline = ExtractionPipelineV2(
converter=mock_converter,
detector=mock_detector,
)
document_id = uuid4()
result = await pipeline.process_document(document_id, "test.xlsx")
assert len(result.records) == 1
record = result.records[0]
# Check content structure
assert "content" in record
assert "text" in record["content"]
assert record["content"]["text"] == "Line 1\nLine 2"
# Check source metadata
assert "_source" in record
assert record["_source"]["document_id"] == str(document_id)
assert record["_source"]["row"] == -1
# Combined raw text uses table_ids (list) instead of table_id
assert record["_source"]["table_ids"] == [0]
class TestPipelineBackwardCompatibility:
"""Test backward compatibility of extract_tables method."""
def test_extract_tables_still_works(self):
"""Test that existing extract_tables method still returns only valid tables."""
html = """
"""
extractor = HtmlTableExtractor()
tables = extractor.extract_tables(html)
# Should only return the valid table
assert len(tables) == 1
assert tables[0].col_count == 2
if __name__ == "__main__":
pytest.main([__file__, "-v"])