#!/usr/bin/env python3 """ Integration test for Large Table Extractor. Tests the full extraction pipeline with a real Excel file. Usage: uv run python tests/integration/test_large_table_integration.py [file_path] """ import sys import time from pathlib import Path from uuid import uuid4 # Add project root to path project_root = Path(__file__).parent.parent.parent sys.path.insert(0, str(project_root)) # Import directly from module to avoid pipeline dependencies from src.extraction_v2.large_table_extractor import ( LargeTableExtractor, LargeTableConfig, LargeTableResult, extract_large_tables ) from src.extraction_v2.detect_border import BorderTableDetector def test_config(): """Test LargeTableConfig defaults.""" print("\n[TEST] LargeTableConfig defaults...") config = LargeTableConfig() assert config.max_rows_for_docling == 1000, "Default max_rows_for_docling should be 1000" assert config.header_sample_rows == 10, "Default header_sample_rows should be 10" assert config.chunk_size == 100, "Default chunk_size should be 100" print(" ✓ Default config values correct") def test_is_large_table(): """Test large table classification.""" print("\n[TEST] Large table classification...") extractor = LargeTableExtractor() # Test large table large_table = {'start_row': 1, 'end_row': 2000} assert extractor.is_large_table(large_table) is True, "2000 rows should be large" # Test small table small_table = {'start_row': 1, 'end_row': 500} assert extractor.is_large_table(small_table) is False, "500 rows should not be large" # Test boundary boundary_table = {'start_row': 1, 'end_row': 1000} assert extractor.is_large_table(boundary_table) is False, "Exactly 1000 should not be large" boundary_table = {'start_row': 1, 'end_row': 1001} assert extractor.is_large_table(boundary_table) is True, "1001 rows should be large" print(" ✓ Large table classification works correctly") def test_border_detector_helpers(): """Test BorderTableDetector helper methods.""" print("\n[TEST] BorderTableDetector helpers...") detector = BorderTableDetector() # Test is_large_table assert detector.is_large_table({'start_row': 1, 'end_row': 2000}) is True assert detector.is_large_table({'start_row': 1, 'end_row': 500}) is False # Test classify_tables tables = [ {'start_row': 1, 'end_row': 500, 'sheet': 'Small'}, {'start_row': 1, 'end_row': 2000, 'sheet': 'Large'}, ] large, normal = detector.classify_tables(tables) assert len(large) == 1, "Should have 1 large table" assert len(normal) == 1, "Should have 1 normal table" assert large[0]['sheet'] == 'Large' assert normal[0]['sheet'] == 'Small' print(" ✓ BorderTableDetector helpers work correctly") def test_pipeline_record_conversion(): """Test conversion to pipeline record format.""" print("\n[TEST] Pipeline record conversion...") extractor = LargeTableExtractor() document_id = uuid4() result = LargeTableResult( success=True, sheet_name="TestSheet", table_range="A1:C10", total_rows=10, rows_processed=2, column_count=3, header_depth=2, headers={"A": "Name", "B": "Value"}, records=[ {"Name": "Row1", "Value": 100}, {"Name": "Row2", "Value": 200} ] ) pipeline_records = extractor.to_pipeline_records(result, document_id) assert len(pipeline_records) == 2, "Should have 2 records" assert pipeline_records[0]['document_id'] == str(document_id) assert pipeline_records[0]['sheet_name'] == "TestSheet" assert pipeline_records[0]['extraction_method'] == 'large_table_chunked' assert pipeline_records[0]['row_index'] == 0 assert pipeline_records[1]['row_index'] == 1 print(" ✓ Pipeline record conversion works correctly") def test_full_extraction(file_path: str): """Test full extraction on a real file.""" print(f"\n[TEST] Full extraction on: {file_path}") if not Path(file_path).exists(): print(f" ⚠️ File not found, skipping: {file_path}") return start_time = time.time() # Create extractor with test config (limit rows for faster testing) config = LargeTableConfig(max_rows_to_process=200) extractor = LargeTableExtractor(config) # Detect tables print(" Detecting tables...") tables = extractor.detect_tables(file_path) print(f" Found {len(tables)} table(s)") if not tables: print(" ⚠️ No tables detected") return # Find large tables large_tables = [t for t in tables if t.get('is_large', False)] print(f" Large tables: {len(large_tables)}") if not large_tables: print(" ⚠️ No large tables found, testing with first table") large_tables = tables[:1] # Extract first large table table = large_tables[0] print(f" Extracting: {table['sheet']} - {table['range']}") result = extractor.extract_table(file_path, table, uuid4()) elapsed = time.time() - start_time if result.success: print(f" ✓ Extraction successful!") print(f" - Rows processed: {result.rows_processed}") print(f" - Columns: {result.column_count}") print(f" - Header depth: {result.header_depth}") print(f" - Processing time: {result.processing_time:.2f}s") print(f" - Total test time: {elapsed:.2f}s") # Show sample headers if result.headers: sample_headers = list(result.headers.items())[:3] print(f" - Sample headers:") for col, name in sample_headers: print(f" {col}: {name[:50]}{'...' if len(name) > 50 else ''}") # Show sample data if result.records: print(f" - Sample data (first row):") first_row = result.records[0] for key, value in list(first_row.items())[:3]: print(f" {key[:30]}: {str(value)[:30]}") else: print(f" ✗ Extraction failed: {result.error}") def run_all_tests(file_path: str = None): """Run all tests.""" print("=" * 60) print(" LARGE TABLE EXTRACTOR - INTEGRATION TESTS") print("=" * 60) # Unit-style tests test_config() test_is_large_table() test_border_detector_helpers() test_pipeline_record_conversion() # Full extraction test if file_path: test_full_extraction(file_path) else: # Try default test file default_file = "CustomerDocument/уВ╖уВЩуГзуГХуВЩф╕Ашжз/1.2 JP1шинхоЪщаЕчЫо.xlsx" if Path(default_file).exists(): test_full_extraction(default_file) else: print("\n[INFO] No test file available for full extraction test") print("\n" + "=" * 60) print(" ALL TESTS PASSED!") print("=" * 60) if __name__ == "__main__": file_path = sys.argv[1] if len(sys.argv) > 1 else None run_all_tests(file_path)