doczyai-pipelines/fieldExtraction/tests/one_to_n_test.py

import unittest
from unittest.mock import patch, MagicMock

import src.investment.one_to_n_funcs as one_to_n_funcs


class TestOneToN(unittest.TestCase):
    def setUp(self):
        self.exhibit_text = "Sample exhibit text with some content"
        self.filename = "test_file.pdf"

    @patch('src.utils.llm_utils.invoke_claude')
    def test_get_reimbursement_primary(self, mock_invoke_claude):
        # Setup
        mock_fields = MagicMock()
        mock_fields.get_prompt_dict.return_value = {"field1": "prompt1"}
        mock_invoke_claude.return_value = '{"answer1": "value1"}'
        
        # Execute
        result = one_to_n_funcs.get_reimbursement_primary(
            mock_fields, self.exhibit_text, self.filename
        )
        
        # Assert
        self.assertEqual(result, {"answer1": "value1"})
        mock_invoke_claude.assert_called_once()

    @patch('src.utils.llm_utils.invoke_claude')
    def test_get_special_cases(self, mock_invoke_claude):
        # Setup
        test_answers = [{
            "SERVICE_TERM": "service1",
            "REIMB_TERM": "reimbursement1"
        }]
        mock_invoke_claude.return_value = "|N/A|"
        
        # Execute
        result = one_to_n_funcs.get_special_cases(test_answers, self.filename)
        
        # Assert
        self.assertEqual(len(result), 1)
        self.assertEqual(result[0]["CARVEOUT_CD"], "N/A")
        self.assertEqual(result[0]["CARVEOUT_IND"], "N")
        self.assertEqual(result[0]["DEFAULT_IND"], "N")

    @patch('src.utils.llm_utils.invoke_claude')
    def test_methodology_breakout_single_row(self, mock_invoke_claude):
        # Setup
        test_answer = {
            "SERVICE_TERM": "service1",
            "REIMB_TERM": "Fee Schedule"
        }
        mock_invoke_claude.side_effect = [
            '{"AARETE_DERIVED_REIMB_METHOD": "Fee Schedule", "REIMB_PCT_RATE": "80%"}',
            '{"FEE_SCHEDULE": "TEST", "AARETE_DERIVED_FEE_SCHEDULE": "TEST", "AARETE_DERIVED_FEE_SCHEDULE_VERSION": "1.0"}'
        ]
        
        # Execute
        result = one_to_n_funcs.methodology_breakout_single_row(
            test_answer,
            {"methodology": "questions"},
            {"fs": "questions"},
            self.filename
        )
        
        # Assert
        self.assertEqual(len(result), 1)
        self.assertEqual(result[0]["AARETE_DERIVED_REIMB_METHOD"], "Fee Schedule")
        self.assertEqual(result[0]["FEE_SCHEDULE"], "TEST")

    def test_combine_one_to_n_answers(self):
        # Setup
        exhibit_level = {"EXHIBIT_TITLE": "Test Exhibit"}
        reimbursement_level = [
            {"REIMB_TERM": "term1"},
            {"REIMB_TERM": "term2"}
        ]
        tin_npi = {"TIN": "123456789"}
        
        # Execute
        result = one_to_n_funcs.combine_one_to_n_answers(
            exhibit_level, reimbursement_level, tin_npi
        )
        
        # Assert
        self.assertEqual(len(result), 2)
        self.assertTrue(all("EXHIBIT_TITLE" in d for d in result))
        self.assertTrue(all("TIN" in d for d in result))
        self.assertTrue(all("REIMB_TERM" in d for d in result))

    @patch('src.investment.one_to_n_funcs.get_reimbursement_primary')
    @patch('src.investment.one_to_n_funcs.get_special_cases')
    @patch('src.investment.one_to_n_funcs.validate_reimbursements_for_llm')
    def test_reimbursement_level(self, mock_validate_llm, mock_special_cases, mock_reimbursement_primary):
        # Setup
        mock_reimbursement_primary.return_value = [{"SERVICE_TERM": "Test Service", "REIMB_TERM": "term1 paid at rate"}]
        mock_validate_llm.return_value = True
        mock_special_cases.return_value = [{"SERVICE_TERM": "Test Service", "REIMB_TERM": "term1 paid at rate", "CARVEOUT_IND": "N"}]

        mock_fields = MagicMock()
        exhibit_page = "1"
        seen_pairs = set()
        exhibit_lesser_of = "N/A"  # Simulate no lesser of statement
        rate_escalator_statement = "N/A"  # Simulate no rate escalator statement
        
        # Execute
        result = one_to_n_funcs.reimbursement_level(
            self.exhibit_text, self.filename, mock_fields, exhibit_page, seen_pairs, exhibit_lesser_of
        )
        
        # Assert
        self.assertEqual(len(result), 1)
        mock_reimbursement_primary.assert_called_once()
        mock_special_cases.assert_called_once()

    @patch('src.investment.one_to_n_funcs.get_reimbursement_primary')
    @patch('src.investment.one_to_n_funcs.validate_reimbursements_for_llm')
    def test_reimbursement_level_deduplication(self, mock_validate_llm, mock_reimbursement_primary):
        """Test that duplicate service-reimbursement pairs are properly filtered"""
        # Setup - simulate getting duplicate pairs
        mock_reimbursement_primary.return_value = [
            {"SERVICE_TERM": "Service A", "REIMB_TERM": "paid at 80% of fee schedule"},
            {"SERVICE_TERM": "Service B", "REIMB_TERM": "$50 per visit"},
            {"SERVICE_TERM": "Service A", "REIMB_TERM": "paid at 80% of fee schedule"}  # Duplicate
        ]
        mock_validate_llm.return_value = True
        
        mock_fields = MagicMock()
        exhibit_page = "1"
        seen_pairs = set()
        
        with patch('src.investment.one_to_n_funcs.get_special_cases') as mock_special_cases:
            
            # Mock the downstream functions
            mock_special_cases.side_effect = lambda x, _: x  # Pass through
            
            # Execute
            result = one_to_n_funcs.reimbursement_level(
                self.exhibit_text, self.filename, mock_fields,
                exhibit_page, seen_pairs, exhibit_lesser_of="N/A"
            )
            
            # Assert - should only have 2 unique pairs
            self.assertEqual(len(result), 2)
            # Verify seen_pairs was updated
            self.assertEqual(len(seen_pairs), 2)
            self.assertIn(("1", "Service A", "paid at 80% of fee schedule"), seen_pairs)
            self.assertIn(("1", "Service B", "$50 per visit"), seen_pairs)

    @patch('src.investment.one_to_n_funcs.get_reimbursement_primary')
    @patch('src.investment.one_to_n_funcs.validate_reimbursements_for_llm')
    def test_reimbursement_level_cross_exhibit_deduplication(self, mock_validate_llm, mock_reimbursement_primary):
        """Test that pairs seen in previous exhibits are filtered out"""
        # Setup - pre-populate seen_pairs as if previous exhibit processed
        seen_pairs = {("2", "Service A", "paid at 80% of fee schedule"),}
        mock_validate_llm.return_value = True
        
        # Current exhibit has same pair plus new one
        mock_reimbursement_primary.return_value = [
            {"SERVICE_TERM": "Service A", "REIMB_TERM": "paid at 80% of fee schedule"},  # Should be filtered
            {"SERVICE_TERM": "Service B", "REIMB_TERM": "$50 per visit"}   # Should be kept
        ]
        
        mock_fields = MagicMock()
        exhibit_page = "2"
        
        with patch('src.investment.one_to_n_funcs.get_special_cases') as mock_special_cases:
            
            mock_special_cases.side_effect = lambda x, _: x
            
            # Execute
            result = one_to_n_funcs.reimbursement_level(
                self.exhibit_text, self.filename, mock_fields,
                exhibit_page, seen_pairs, exhibit_lesser_of="N/A"
            )
            
            # Assert - should only process the new pair
            self.assertEqual(len(result), 1)
            self.assertEqual(result[0]["SERVICE_TERM"], "Service B")
            # Verify both pairs are now in seen_pairs
            self.assertEqual(len(seen_pairs), 2)

    @patch('src.investment.one_to_n_funcs.get_reimbursement_primary')
    @patch('src.investment.one_to_n_funcs.validate_reimbursements_for_llm')
    def test_reimbursement_level_cross_exhibit_deduplication_same_base_page(self, mock_validate_llm, mock_reimbursement_primary):
        """Test that pairs seen in previous exhibits within same base page are filtered out"""
        # Setup - pre-populate seen_pairs as if previous exhibit in same base page processed
        seen_pairs = {("2", "Service A", "paid at 80% of fee schedule"),}
        mock_validate_llm.return_value = True
        
        # Current exhibit (2.1) has same pair plus new one - same base page
        mock_reimbursement_primary.return_value = [
            {"SERVICE_TERM": "Service A", "REIMB_TERM": "paid at 80% of fee schedule"},  # Should be filtered (duplicate in base page 2)
            {"SERVICE_TERM": "Service B", "REIMB_TERM": "$50 per visit"}   # Should be kept
        ]
        
        mock_fields = MagicMock()
        mock_dataset = {}
        exhibit_page = "2.1"  # Same base page (2) as seen_pairs
        
        with patch('src.investment.one_to_n_funcs.get_special_cases') as mock_special_cases:
            
            mock_special_cases.side_effect = lambda x, _: x
            
            # Execute
            result = one_to_n_funcs.reimbursement_level(
                self.exhibit_text, self.filename, mock_fields, 
                exhibit_page, seen_pairs, exhibit_lesser_of="N/A"
            )
            
            # Assert - should only process the new pair (Service A filtered out)
            self.assertEqual(len(result), 1)
            self.assertEqual(result[0]["SERVICE_TERM"], "Service B")
            # Verify both pairs are now in seen_pairs
            self.assertEqual(len(seen_pairs), 2)
            self.assertIn(("2", "Service B", "$50 per visit"), seen_pairs)

    @patch('src.investment.one_to_n_funcs.get_reimbursement_primary')
    @patch('src.investment.one_to_n_funcs.validate_reimbursements_for_llm')
    def test_reimbursement_level_cross_exhibit_deduplication_different_base_page(self, mock_validate_llm, mock_reimbursement_primary):
        """Test that pairs from different base pages are NOT filtered out"""
        # Setup - pre-populate seen_pairs as if previous exhibit from different base page processed
        seen_pairs = {("23", "Service A", "paid at 80% of fee schedule"),}
        mock_validate_llm.return_value = True
        
        # Current exhibit has same pair but from different base page
        mock_reimbursement_primary.return_value = [
            {"SERVICE_TERM": "Service A", "REIMB_TERM": "paid at 80% of fee schedule"},  # Should be kept (different base page)
            {"SERVICE_TERM": "Service B", "REIMB_TERM": "$50 per visit"}   # Should be kept
        ]
        
        mock_fields = MagicMock()
        mock_dataset = {}
        exhibit_page = "24.0"  # Different base page (24) from seen_pairs (23)
        
        with patch('src.investment.one_to_n_funcs.get_special_cases') as mock_special_cases:
            
            mock_special_cases.side_effect = lambda x, _: x
            
            # Execute
            result = one_to_n_funcs.reimbursement_level(
                self.exhibit_text, self.filename, mock_fields, 
                exhibit_page, seen_pairs, exhibit_lesser_of="N/A"
            )
            
            # Assert - should process both pairs (no filtering due to different base page)
            self.assertEqual(len(result), 2)
            # Verify all pairs are now in seen_pairs (original + 2 new)
            self.assertEqual(len(seen_pairs), 3)
            self.assertIn(("23", "Service A", "paid at 80% of fee schedule"), seen_pairs)  # Original
            self.assertIn(("24", "Service A", "paid at 80% of fee schedule"), seen_pairs)  # New
            self.assertIn(("24", "Service B", "$50 per visit"), seen_pairs)  # New

    def test_deduplicate_with_accumulator_base_page_extraction(self):
        """Test that base page extraction works correctly for various exhibit page formats"""
        answers = [
            {"SERVICE_TERM": "Service A", "REIMB_TERM": "term1"},
            {"SERVICE_TERM": "Service B", "REIMB_TERM": "term2"}
        ]
        
        test_cases = [
            ("23", "23"),      # No decimal
            ("23.0", "23"),    # Single decimal  
            ("23.1", "23"),    # Single decimal with sub-exhibit
            ("23.0.0", "23"),  # Double decimal
            ("123.45.67", "123")  # Multiple decimals
        ]
        
        for exhibit_page, expected_base_page in test_cases:
            seen_pairs = set()
            
            result = one_to_n_funcs.deduplicate_with_accumulator(
                answers, seen_pairs, exhibit_page, "test_file.pdf"
            )
            
            # Verify correct base page was used in the keys
            expected_keys = {
                (expected_base_page, "Service A", "term1"),
                (expected_base_page, "Service B", "term2")
            }
            self.assertEqual(seen_pairs, expected_keys, f"Failed for exhibit_page: {exhibit_page}")

    def test_deduplicate_with_accumulator_within_vs_across_base_pages(self):
        """Test deduplication behavior within same vs different base pages"""
        answers = [
            {"SERVICE_TERM": "Service A", "REIMB_TERM": "term1"},
            {"SERVICE_TERM": "Service B", "REIMB_TERM": "term2"}
        ]
        
        # First exhibit in base page 23
        seen_pairs = set()
        result1 = one_to_n_funcs.deduplicate_with_accumulator(
            answers, seen_pairs, "23.0", "test_file.pdf"
        )
        self.assertEqual(len(result1), 2)  # Both kept
        self.assertEqual(len(seen_pairs), 2)
        
        # Second exhibit in same base page 23 - should filter duplicates
        result2 = one_to_n_funcs.deduplicate_with_accumulator(
            answers, seen_pairs, "23.1", "test_file.pdf"
        )
        self.assertEqual(len(result2), 0)  # Both filtered (duplicates within base page)
        self.assertEqual(len(seen_pairs), 2)  # No new pairs added
        
        # Third exhibit in different base page 24 - should keep all
        result3 = one_to_n_funcs.deduplicate_with_accumulator(
            answers, seen_pairs, "24.0", "test_file.pdf"
        )
        self.assertEqual(len(result3), 2)  # Both kept (different base page)
        self.assertEqual(len(seen_pairs), 4)  # 2 new pairs added
        
        # Verify final state
        expected_pairs = {
            ("23", "Service A", "term1"), ("23", "Service B", "term2"),
            ("24", "Service A", "term1"), ("24", "Service B", "term2")
        }
        self.assertEqual(seen_pairs, expected_pairs)

    @patch('src.investment.one_to_n_funcs.validate_reimbursements_for_llm')
    def test_reimbursement_level_all_duplicates(self, mock_validate_llm):
        """Test behavior when all pairs are duplicates"""
        seen_pairs = {("3", "Service A", "paid at 80% of fee schedule"), ("3", "Service B", "$50 per visit")}
        mock_validate_llm.return_value = True
        with patch('src.investment.one_to_n_funcs.get_reimbursement_primary') as mock_primary:
            mock_primary.return_value = [
                {"SERVICE_TERM": "Service A", "REIMB_TERM": "paid at 80% of fee schedule"},
                {"SERVICE_TERM": "Service B", "REIMB_TERM": "$50 per visit"}
            ]
            
            mock_fields = MagicMock()
            mock_dataset = {}
            exhibit_page = "3"
            
            # Execute
            result = one_to_n_funcs.reimbursement_level(
                self.exhibit_text, self.filename, mock_fields,
                exhibit_page, seen_pairs, exhibit_lesser_of="N/A"
            )
            
            # Assert - should return empty list, no downstream processing
            self.assertEqual(result, [])

    def test_deduplicate_with_accumulator(self):
        """Test the deduplication function directly"""
        answers = [
            {"SERVICE_TERM": "Service A", "REIMB_TERM": "paid at 80% of fee schedule"},
            {"SERVICE_TERM": "Service B", "REIMB_TERM": "$50 per visit"},
            {"SERVICE_TERM": "Service A", "REIMB_TERM": "paid at 80% of fee schedule"},  # Duplicate
            {"SERVICE_TERM": " Service C ", "REIMB_TERM": " $150 "}  # Test stripping
        ]
        
        seen_pairs = set()
        
        result = one_to_n_funcs.deduplicate_with_accumulator(
            answers, seen_pairs, "test_page", "test_file.pdf"
        )
        
        # Should have 3 unique pairs (duplicates removed)
        self.assertEqual(len(result), 3)
        self.assertEqual(len(seen_pairs), 3)
        
        # Verify stripping works
        self.assertIn(("test_page", "Service C", "$150"), seen_pairs)

if __name__ == '__main__':
    unittest.main()