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