doczyai-pipelines/fieldExtraction/tests/preprocessing_funcs_test.py

import pytest
from src.preprocessing_funcs import (
    remove_page_indicators,
    clean_law_symbols,
    filter_quick_review,
    get_exhibit_pages,
)
import src.prompts.investment_prompts as investment_prompts
import src.config as config

from src.client.smart_chunking_funcs import smart_chunk_ac
from unittest.mock import patch


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.client.smart_chunking_funcs.chunk_{chunking_fn_name}", return_value=mock_return)   
        mock_keyword_search = mocker.patch("src.client.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


    @patch('src.utils.llm_utils.invoke_claude')
    @patch('src.utils.string_utils.extract_text_from_delimiters')
    def test_get_exhibit_pages(self, mock_extract_text, mock_invoke_claude):
        """Test exhibit page detection with both base and decimal pages"""
        # Setup test data
        text_dict = {
            "1": "First page content",
            "2.0": "Second page content with exhibit header",
            "2.1": "Continuation of second page",
            "2.2": "More continuation of second page",
            "3.0": "Third page content with exhibit header",
            "4": "Regular page content"
        }
        
        filename = "test_contract.pdf"
        
        # Configure mocks
        mock_invoke_claude.return_value = "raw_llm_response"
        mock_extract_text.side_effect = [
            "Exhibit A",     # For page 2.0
            "N/A",          # For page 4
            "Exhibit B"      # For page 3.0
        ]
        
        # Execute
        exhibit_pages, exhibit_headers = get_exhibit_pages(text_dict, filename)
        
        # Assert
        expected_exhibit_pages = ["1", "2.0", "2.1", "2.2", '4']
        
        assert exhibit_pages == expected_exhibit_pages
        
        # Verify LLM was only called for base pages
        assert mock_invoke_claude.call_count == 3
        
        # Verify prompt content and model
        mock_invoke_claude.assert_any_call(
            investment_prompts.EXHIBIT_HEADER("Second page content with exhibit header"[0:400]),
            config.MODEL_ID_CLAUDE3_HAIKU,
            filename,
            max_tokens=150
        )