doczyai-pipelines/fieldExtraction/tests/preprocessing_funcs_test.py

import pytest
from src.preprocessing_funcs import (
    clean_newlines,
    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 single newlines
        ("This is a test.\nThis is only a test.", "This is a test. This is only a test."),
        # 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."),
        # Handles mixed newlines and page numbers
        ("Page 1 of 10\n\nLine 1\nLine 2\nPage 2 of 10\n\nLine 3", " Line 1 Line 2  Line 3"),
    ])
    def test_clean_newlines(self, input_text, expected_output):
        assert clean_newlines(input_text) == expected_output

    # Test cases for split_text
    @pytest.mark.parametrize("input_text, expected_output", [
        # Splits on page markers
        (
            "Title Start of Page No. = 1\nContent1\nStart of Page No. = 2\nContent2",
            {"Title" : "Title ", "1": "Content1\n", "2": "Content2"}
        ),
        # Ignores "Document" key
        (
            "Document\nStart of Page No. = 1\nContent1",
            {"1": "Content1"}
        ),
        # Handles empty input
        ("", {}),
        # Handles no page markers
        ("This is a single page.", {"This": "This is a single page."}),
        # Handles multiple page markers
        (
            "Introduction\nStart of Page No. = 1\nContent1\nStart of Page No. = 2\nContent2\nStart of Page No. = 3\nContent3",
            {'Introduction': 'Introduction\n', '1': 'Content1\n', '2': 'Content2\n', '3': 'Content3'}
        ),
    ])
    def test_split_text(self, input_text, expected_output):
        assert split_text(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)