import pytest from src.preprocessing_funcs import split_text from src.preprocessing_funcs import ( remove_page_indicators, split_text, clean_law_symbols, smart_chunk_ac, filter_quick_review, get_exhibit_pages, ) 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 smart_chunk_ac @pytest.mark.parametrize("methodology, mock_return, chunking_fn_name, expected_chunk, expected_pages", [ # OR methodology ("or", [1, 2], "or", "Content1\nContent2", {"test": [1, 2]}), # Hierarchy methodology ("hierarchy", [1, 2], "hierarchical", "Content1\nContent2", {"test": [1, 2]}), # AND methodology ("and", [1], "and", "Content1", {"test": [1]}), # Regex methodology ("regex", [1, 2], "regex", "Content1\nContent2", {"test": [1, 2]}), # Include/Exclude methodology ("include_exclude", [1], "with_include_exclude_keywords", "Content1", {"test": [1]}), ]) def test_smart_chunk_ac(self, mocker, methodology, mock_return, chunking_fn_name, expected_chunk, expected_pages): # Mock the chunking functions mock_chunk_func = mocker.patch(f"src.smart_chunking_funcs.chunk_{chunking_fn_name}", return_value=mock_return) mock_keyword_search = mocker.patch("src.smart_chunking_funcs.keyword_search", return_value={"test": mock_return}) text_dict = {"1": "Content1", "2": "Content2"} keyword_mappings = {"group1": {"methodology": methodology, "keywords": ["test"], "case_sensitive": False}} if methodology == "regex": keyword_mappings["group1"]["regex"] = "Content" if methodology == "include_exclude": del keyword_mappings["group1"]["keywords"] keyword_mappings["group1"]["included_keywords"] = ["Content"] keyword_mappings["group1"]["excluded_keywords"] = [] result = smart_chunk_ac(text_dict, "Content1\nContent2", keyword_mappings) assert result["group1"] == expected_chunk assert result["group1_pages"] == expected_pages mock_chunk_func.assert_called_once() mock_keyword_search.assert_called_once() # 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 get_exhibit_pages @pytest.mark.parametrize("mock_response, expected_pages, text_dict", [ # Exhibit page detected (["|Y|","|N|"], ["1"], {"1": "Exhibit content", "2": "Non-exhibit content"}), # No exhibit page detected ([], [], {}), # Multiple exhibit pages (["|Y|","|Y|"], ["1", "2"], {"1": "Exhibit content", "2": "Attachment"}), # Invalid response (["Invalid"], [], {"1": "Non-exhibit content"}), ]) def test_get_exhibit_pages(self, mocker, mock_response, expected_pages, text_dict): # Mock the LLM response mocker_invoke_claude = mocker.patch("src.utils.llm_utils.invoke_claude", side_effect=mock_response) filename = "test.pdf" assert get_exhibit_pages(text_dict, filename) == expected_pages # # verify calls made # assert mocker_invoke_claude.call_count == len(text_dict)