doczyai-pipelines/fieldExtraction/tests/test_preprocessing_funcs.py

import pytest
from src.investment.preprocessing_funcs import (chunk_by_exhibit,
                                                clean_law_symbols,
                                                filter_quick_review,
                                                remove_page_indicators,
                                                simplify_exhibit)


class TestPreprocessingFuncs:
    # Test cases for clean_newlines
    @pytest.mark.parametrize(
        "input_text, expected_output",
        [
            # Removes page number indicators
            ("Page 1 of 10\n\nThis is a test.", " This is a test."),
            # Handles multiple newlines
            ("Line 1\n\nLine 2\n\nLine 3", "Line 1\n\nLine 2\n\nLine 3"),
            # Handles no newlines
            ("This is a single line.", "This is a single line."),
        ],
    )
    def test_clean_newlines(self, input_text, expected_output):
        assert remove_page_indicators(input_text) == expected_output

    # Test cases for clean_law_symbols
    @pytest.mark.parametrize(
        "input_text, expected_output",
        [
            # Replaces double dollar signs
            ("This is a $$ test.", "This is a $ test."),
            # Replaces U.S.C. symbol
            ("U.S.C. $1234", "U.S.C.§1234"),
            # Replaces C.F.R. symbol
            ("C.F.R. $1234", "C.F.R.§1234"),
            # Replaces $ followed by three decimal places
            ("$123.456", "§123.456"),
            # Handles no replacements
            ("This is a normal text.", "This is a normal text."),
            # Handles mixed replacements
            (
                "U.S.C. $1234 and C.F.R. $5678 and $123.456",
                "U.S.C.§1234 and C.F.R.§5678 and §123.456",
            ),
        ],
    )
    def test_clean_law_symbols(self, input_text, expected_output):
        assert clean_law_symbols(input_text) == expected_output

    # Test cases for filter_quick_review
    @pytest.mark.parametrize(
        "input_dict, expected_contract, expected_quick_review",
        [
            # Filters quick review pages
            (
                {"1": "QUICK REVIEW", "2": "Normal content", "3": "COVER SHEET"},
                {"2": "Normal content"},
                {"1": "QUICK REVIEW", "3": "COVER SHEET"},
            ),
            # No quick review pages
            (
                {"1": "Normal content", "2": "More content"},
                {"1": "Normal content", "2": "More content"},
                {},
            ),
            # All quick review pages
            (
                {"1": "QUICK REVIEW", "2": "TOP SHEET"},
                {},
                {"1": "QUICK REVIEW", "2": "TOP SHEET"},
            ),
            # Case-insensitive matching
            (
                {"1": "quick review", "2": "Normal content"},
                {"2": "Normal content"},
                {"1": "quick review"},
            ),
        ],
    )
    def test_filter_quick_review(
        self, input_dict, expected_contract, expected_quick_review
    ):
        contract, quick_review = filter_quick_review(input_dict)
        assert contract == expected_contract
        assert quick_review == expected_quick_review

    # Test cases for chunk_by_exhibit
    def test_chunk_by_exhibit_basic(self):
        """Test basic exhibit chunking functionality."""
        text_dict = {"1": "Page 1", "2": "Page 2", "3": "Page 3", "4": "Page 4"}
        exhibit_pages = ["1", "3"]
        
        result = chunk_by_exhibit(text_dict, exhibit_pages)
        
        assert result == {"1": ["1", "2"], "3": ["3", "4"]}

    def test_chunk_by_exhibit_single_exhibit(self):
        """Test chunking when all pages belong to one exhibit."""
        text_dict = {"1": "Page 1", "2": "Page 2", "3": "Page 3"}
        exhibit_pages = ["1"]
        
        result = chunk_by_exhibit(text_dict, exhibit_pages)
        
        assert result == {"1": ["1", "2", "3"]}

    def test_chunk_by_exhibit_empty_exhibit_pages(self):
        """Test chunking when exhibit_pages is empty - each page becomes its own exhibit."""
        text_dict = {"1": "Page 1", "2": "Page 2", "3": "Page 3"}
        exhibit_pages = []
        
        result = chunk_by_exhibit(text_dict, exhibit_pages)
        
        assert result == {"1": ["1"], "2": ["2"], "3": ["3"]}

    def test_chunk_by_exhibit_exhibit_at_start(self):
        """Test chunking when exhibit starts at first page."""
        text_dict = {"1": "Page 1", "2": "Page 2", "3": "Page 3"}
        exhibit_pages = ["1", "2"]
        
        result = chunk_by_exhibit(text_dict, exhibit_pages)
        
        assert result == {"1": ["1"], "2": ["2", "3"]}

    def test_chunk_by_exhibit_multiple_exhibits(self):
        """Test chunking with multiple exhibits."""
        text_dict = {"1": "Page 1", "2": "Page 2", "3": "Page 3", "4": "Page 4", "5": "Page 5"}
        exhibit_pages = ["1", "3", "5"]
        
        result = chunk_by_exhibit(text_dict, exhibit_pages)
        
        assert result == {"1": ["1", "2"], "3": ["3", "4"], "5": ["5"]}

    # Test cases for simplify_exhibit
    def test_simplify_exhibit_no_tables(self):
        """Test simplify_exhibit when there are no tables."""
        text_dict = {"1": "Page 1 content", "2": "Page 2 content"}
        exhibit_page_nums = ["1", "2"]
        current_page_num = "1"
        
        result = simplify_exhibit(text_dict, exhibit_page_nums, current_page_num)
        
        assert result == "Page 1 content\nPage 2 content"

    def test_simplify_exhibit_with_tables_current_page_preserved(self):
        """Test that current page tables are preserved while others are simplified."""
        text_dict = {
            "1": "Before table-------Table Start--------Table content-------Table End--------After",
            "2": "Before table-------Table Start--------Table content-------Table End--------After"
        }
        exhibit_page_nums = ["1", "2"]
        current_page_num = "1"
        
        result = simplify_exhibit(text_dict, exhibit_page_nums, current_page_num)
        
        # Current page (1) should have full table content
        assert "-------Table Start--------" in result.split("\n")[0]
        # Other page (2) should have table replaced with placeholder
        # Note: Table counter is incremented when counting tables on current page, so page 2 gets "Table 2"
        assert "Table 2" in result.split("\n")[1]
        assert "-------Table Start--------" not in result.split("\n")[1]

    def test_simplify_exhibit_table_numbering_consistency(self):
        """Test that table numbering is consistent across pages."""
        text_dict = {
            "1": "Text-------Table Start--------Table 1-------Table End--------Text",
            "2": "Text-------Table Start--------Table 2-------Table End--------Text",
            "3": "Text-------Table Start--------Table 3-------Table End--------Text"
        }
        exhibit_page_nums = ["1", "2", "3"]
        current_page_num = "2"
        
        result = simplify_exhibit(text_dict, exhibit_page_nums, current_page_num)
        result_pages = result.split("\n")
        
        # Page 1 (simplified) should have "Table 1"
        assert "Table 1" in result_pages[0]
        # Page 2 (current, preserved) should have full table content
        assert "-------Table Start--------" in result_pages[1]
        # Page 3 (simplified) should have "Table 3" 
        # (counter is 2 after page 1, then incremented to 3 when counting page 2's table)
        assert "Table 3" in result_pages[2]

    def test_simplify_exhibit_multiple_tables_per_page(self):
        """Test simplify_exhibit with multiple tables on a page."""
        text_dict = {
            "1": "Text-------Table Start--------T1-------Table End--------Text-------Table Start--------T2-------Table End--------",
            "2": "Text-------Table Start--------T3-------Table End--------"
        }
        exhibit_page_nums = ["1", "2"]
        current_page_num = "2"
        
        result = simplify_exhibit(text_dict, exhibit_page_nums, current_page_num)
        result_pages = result.split("\n")
        
        # Page 1 should have both tables replaced: "Table 1" and "Table 2"
        assert "Table 1" in result_pages[0]
        assert "Table 2" in result_pages[0]
        # Page 2 (current) should preserve table content
        assert "-------Table Start--------" in result_pages[1]

    def test_simplify_exhibit_missing_page_in_text_dict(self):
        """Test simplify_exhibit when a page number is missing from text_dict."""
        text_dict = {"1": "Page 1 content"}
        exhibit_page_nums = ["1", "2"]
        current_page_num = "1"
        
        result = simplify_exhibit(text_dict, exhibit_page_nums, current_page_num)
        
        # Should handle missing page gracefully (empty string)
        assert "Page 1 content" in result
        assert result.count("\n") == 1  # Two pages separated by newline

    def test_simplify_exhibit_table_without_end_marker(self):
        """Test simplify_exhibit when table has no end marker."""
        text_dict = {
            "1": "Text-------Table Start--------Table content without end",
            "2": "Page 2 content"
        }
        exhibit_page_nums = ["1", "2"]
        current_page_num = "2"
        
        result = simplify_exhibit(text_dict, exhibit_page_nums, current_page_num)
        
        # Should still process and replace table, stopping at end of text
        assert "Table 1" in result or "-------Table Start--------" in result