doczyai-pipelines/fieldExtraction/tests/preprocessing_integration_test.py

import unittest
from unittest.mock import patch

import pytest

from src.investment.preprocess import clean_tables


class TestPreprocessingIntegration(unittest.TestCase):

    def test_clean_tables_simple_end_to_end(self):
        """Test the complete simplified table cleaning pipeline."""
        text_dict = {
            "1": "Metadata\n-------Table Start--------\nMain Table\n[['A', 'B']] + [[f'Row{i}%', str(i)] for i in range(1, 8)]\n-------Table End--------",
            "2": "-------Table Start--------\n\n[['Row8', '8%'], ['Row9', '9%']]\n-------Table End--------",  # Continuation
            "3": "Text without tables",
        }

        # This should: 1) Combine continuations, 2) Split by row count
        result = clean_tables(text_dict, "test_file.txt")

        # Should have combined page 2 into page 1, then split the long table
        self.assertNotIn("2", result)  # Continuation removed
        self.assertIn("3", result)  # Non-table page preserved
        # Should have split pages like "1.0", "1.1", etc.

    def test_clean_tables_small_table_gets_dot_zero_suffix(self):
        """Test that pages with small tables get .0 suffix but aren't split."""
        text_dict = {
            "15": "Rate Table\n-------Table Start--------\nRate Table\n[['HCPC', 'Rate'], ['T1019', '$5.00']]\n-------Table End--------"
        }

        result = clean_tables(text_dict, "test_file.pdf")

        # Should get .0 suffix because it contains tables
        self.assertIn("15.0", result)
        self.assertNotIn("15", result)

        # Content should be preserved
        self.assertIn("T1019", result["15.0"])

        # Should not be split (small table)
        chunk_keys = [k for k in result.keys() if k.startswith("15.")]
        self.assertEqual(len(chunk_keys), 1)

    def test_clean_tables_continuation_gets_combined_and_dot_zero(self):
        """Test that continuations get combined and result gets .0 suffix."""
        text_dict = {
            "1": "Rate Table\n-------Table Start--------\nRate Table Header\n[['HCPC', 'Rate'], ['T1019', '$5.00']]\n-------Table End--------",
            "2": "-------Table Start--------\n\n[['T1021', '$4.00']]\n-------Table End--------",  # Continuation
        }

        result = clean_tables(text_dict, "test_file.pdf")

        # Should combine and get .0 suffix since contains tables
        self.assertIn("1.0", result)
        self.assertNotIn("1", result)
        self.assertNotIn("2", result)  # Absorbed into Page 1

        # Verify combined content
        self.assertIn("T1019", result["1.0"])
        self.assertIn("T1021", result["1.0"])

    def test_clean_tables_large_table_row_based_splitting(self):
        """Test that large tables get split based on row count."""
        # Create a table with many rows that will exceed TABLE_ROW_LIMIT
        large_table_rows = [["HCPC", "Rate"]] + [
            [f"T{i:04d}", f"${i}.00"] for i in range(60)
        ]  # 60 data rows
        text_dict = {
            "31": f"EXHIBIT F\n-------Table Start--------\nPCAP Rates\n{large_table_rows}\n-------Table End--------\nPost-table text"
        }

        result = clean_tables(text_dict, "test_file.pdf")

        # Should create multiple chunks due to row limit (default 50)
        chunk_keys = [k for k in result.keys() if k.startswith("31.")]
        self.assertGreater(
            len(chunk_keys),
            1,
            f"Expected multiple chunks for 60 rows, got: {chunk_keys}",
        )

        # Original page should be gone
        self.assertNotIn("31", result)

        # All table data should be preserved somewhere
        all_content = "".join(result.values())
        self.assertIn("T0000", all_content)
        self.assertIn("T0059", all_content)
        self.assertIn("EXHIBIT F", all_content)

        # Post-table text should be preserved on last chunk
        last_chunk = max(chunk_keys)
        self.assertIn("Post-table text", result[last_chunk])

    def test_clean_tables_no_tables_unchanged(self):
        """Test that pages without tables remain unchanged."""
        text_dict = {
            "5": "This is just text content with no tables.",
            "6": "More text content without any table markers.",
        }

        result = clean_tables(text_dict, "test_file.pdf")

        # Should remain unchanged (no .0 suffix)
        assert "5" in result
        assert "6" in result
        assert "5.0" not in result
        assert "6.0" not in result

        # Content should be identical
        assert result["5"] == text_dict["5"]
        assert result["6"] == text_dict["6"]

    def test_clean_tables_multiple_tables_same_page(self):
        """Test that multiple tables on same page get .0 suffix but stay together."""
        text_dict = {
            "1": "Page metadata\n-------Table Start--------\nFirst Table\n[['Col1', 'Col2'], ['A%', 'B%']]\n-------Table End--------\n\nSome text\n-------Table Start--------\nSecond Table\n[['X', 'Y'], ['1%', '2%']]\n-------Table End--------"
        }

        result = clean_tables(text_dict, "test_file.pdf")

        # Should get .0 suffix because it contains tables
        assert "1.0" in result
        assert "1" not in result

        # Both tables should be present in single chunk
        assert "First Table" in result["1.0"]
        assert "Second Table" in result["1.0"]