doczyai-pipelines/fieldExtraction/tests/preprocessing_test.py

import pytest
from src.investment.preprocess import get_pages_to_process


class TestGetPagesToProcess:
    """Test cases for get_pages_to_process function."""

    def test_regular_pages_only(self):
        """Test processing exhibits with no subpages - should return exhibit pages for chunk processing."""
        reimbursement_exhibits = ["1", "5"]
        exhibit_chunk_mapping = {
            "1": "1",
            "2": "1", 
            "3": "1",
            "5": "5",
            "6": "5",
            "7": "7"
        }
        
        result = get_pages_to_process(reimbursement_exhibits, exhibit_chunk_mapping)
        
        assert result == ["1", "5"]

    def test_subpages_only(self):
        """Test processing exhibits with subpages - should return all subpages individually."""
        reimbursement_exhibits = ["5"]
        exhibit_chunk_mapping = {
            "1": "1",
            "5.1": "5",
            "5.2": "5", 
            "5.3": "5",
            "7": "7"
        }
        
        result = get_pages_to_process(reimbursement_exhibits, exhibit_chunk_mapping)
        
        assert set(result) == {"5.1", "5.2", "5.3"}

    def test_mixed_regular_and_subpages(self):
        """Test processing mix of regular exhibits and exhibits with subpages."""
        reimbursement_exhibits = ["1", "5", "8"]
        exhibit_chunk_mapping = {
            "1": "1",           # Regular exhibit (no subpages)
            "2": "1",
            "5.1": "5",         # Exhibit with subpages
            "5.2": "5",
            "5.3": "5",
            "8": "8",           # Regular exhibit
            "9": "8",
            "10": "8"
        }
        
        result = get_pages_to_process(reimbursement_exhibits, exhibit_chunk_mapping)
        
        expected = ["1", "5.1", "5.2", "5.3", "8"]
        assert set(result) == set(expected)

    def test_exhibit_with_base_page_and_subpages(self):
        """Test exhibit that has both base page and subpages (e.g., '5', '5.1', '5.2')."""
        reimbursement_exhibits = ["5"]
        exhibit_chunk_mapping = {
            "1": "1",
            "5": "5",           # Base page
            "5.1": "5",         # Subpages
            "5.2": "5",
            "7": "7"
        }
        
        result = get_pages_to_process(reimbursement_exhibits, exhibit_chunk_mapping)
        
        # Should return all pages in the exhibit since it has subpages
        assert set(result) == {"5", "5.1", "5.2"}

    def test_empty_reimbursement_exhibits(self):
        """Test handling of empty reimbursement exhibits list."""
        reimbursement_exhibits = []
        exhibit_chunk_mapping = {
            "1": "1",
            "2": "1",
            "5.1": "5",
            "5.2": "5"
        }
        
        result = get_pages_to_process(reimbursement_exhibits, exhibit_chunk_mapping)
        
        assert result == []

    def test_complex_subpage_numbering(self):
        """Test complex subpage numbering patterns."""
        reimbursement_exhibits = ["5", "10"]
        exhibit_chunk_mapping = {
            "1": "1",
            "5.0": "5",         # Different subpage patterns
            "5.1": "5",
            "5.10": "5",        # Double digit subpage
            "5.2": "5",
            "10.1": "10",
            "10.11": "10",      # Double digit subpage
            "10.2": "10"
        }
        
        result = get_pages_to_process(reimbursement_exhibits, exhibit_chunk_mapping)
        
        expected_5 = {"5.0", "5.1", "5.10", "5.2"}
        expected_10 = {"10.1", "10.11", "10.2"}
        
        assert set(result) == expected_5.union(expected_10)

    def test_single_exhibit_single_page(self):
        """Test single exhibit with single page."""
        reimbursement_exhibits = ["1"]
        exhibit_chunk_mapping = {
            "1": "1"
        }
        
        result = get_pages_to_process(reimbursement_exhibits, exhibit_chunk_mapping)
        
        assert result == ["1"]

    def test_decimal_in_exhibit_name_not_subpage(self):
        """Test edge case where exhibit name contains decimal but isn't a subpage."""
        # This might be an edge case depending on your data format
        reimbursement_exhibits = ["1.5"]  # Exhibit named "1.5"
        exhibit_chunk_mapping = {
            "1.5": "1.5",       # This is the exhibit itself, not a subpage
            "2": "1.5"          # Regular page belonging to exhibit "1.5"
        }
        
        result = get_pages_to_process(reimbursement_exhibits, exhibit_chunk_mapping)
        
        # Since "1.5" contains a decimal, it should process pages individually
        assert set(result) == {"1.5", "2"}

    def test_order_preservation(self):
        """Test that the order of processing is preserved/predictable."""
        reimbursement_exhibits = ["1", "5", "8"]
        exhibit_chunk_mapping = {
            "1": "1",
            "5.2": "5",
            "5.1": "5",         # Intentionally out of order
            "8": "8"
        }
        
        result = get_pages_to_process(reimbursement_exhibits, exhibit_chunk_mapping)
        
        # Should maintain order of exhibits in reimbursement_exhibits
        assert result[0] == "1"  # First exhibit
        assert "5.2" in result and "5.1" in result  # Subpages from second exhibit
        assert result[-1] == "8"  # Last exhibit

    def test_multiple_exhibits_with_subpages(self):
        """Test multiple exhibits that all have subpages."""
        reimbursement_exhibits = ["3", "7", "12"]
        exhibit_chunk_mapping = {
            "1": "1",
            "3.1": "3",
            "3.2": "3",
            "7.1": "7",
            "7.2": "7",
            "7.3": "7",
            "12.1": "12",
            "12.2": "12"
        }
        
        result = get_pages_to_process(reimbursement_exhibits, exhibit_chunk_mapping)
        
        expected = ["3.1", "3.2", "7.1", "7.2", "7.3", "12.1", "12.2"]
        assert set(result) == set(expected)


# Integration test with actual preprocessing flow
class TestGetPagesToProcessIntegration:
    """Integration tests that simulate real preprocessing scenarios."""
    
    def test_table_splitting_scenario(self):
        """Test scenario where tables are split creating subpages."""
        # Simulate a scenario where exhibit 5 had a large table that got split
        reimbursement_exhibits = ["1", "5"]
        exhibit_chunk_mapping = {
            "1": "1",           # Regular exhibit
            "2": "1", 
            "5": "5",           # Original exhibit page
            "5.1": "5",         # Table continuation pages
            "5.2": "5",
            "5.3": "5",
            "6": "6"            # Next exhibit (not reimbursement)
        }
        
        result = get_pages_to_process(reimbursement_exhibits, exhibit_chunk_mapping)
        
        # Exhibit 1: regular processing (chunk)
        # Exhibit 5: individual processing (has subpages from table splitting)
        expected = ["1", "5", "5.1", "5.2", "5.3"]
        assert set(result) == set(expected)