doczyai-pipelines/fieldExtraction/tests/dynamic_test.py

import unittest
from unittest.mock import patch, MagicMock

import src.investment.one_to_n_funcs as one_to_n_funcs
from src.prompts.investment_prompts import FieldSet, Field

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')
    @patch('src.prompts.investment_prompts.FieldSet')
    def test_get_exhibit_level_answers(self, mock_fieldset, mock_invoke_claude):
        # Setup
        mock_fieldset_instance = MagicMock()
        mock_fieldset_instance.contains_fields.return_value = True
        mock_fieldset.return_value = mock_fieldset_instance
        
        mock_invoke_claude.return_value = '{"field1": "value1", "field2": "value2"}'
        
        # Execute
        result = one_to_n_funcs.get_exhibit_level_answers(self.exhibit_text, self.filename)
        
        # Assert
        self.assertEqual(result, {"field1": "value1", "field2": "value2"})
        mock_invoke_claude.assert_called_once()

    @patch('src.utils.llm_utils.invoke_claude')
    def test_get_exhibit_header(self, mock_invoke_claude):
        # Setup
        mock_invoke_claude.return_value = "|Sample Header|"
        
        # Execute
        result = one_to_n_funcs.get_exhibit_header(self.exhibit_text, self.filename)
        
        # Assert
        self.assertEqual(result, "Sample Header")
        mock_invoke_claude.assert_called_once()

    @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.get_methodology_breakout')
    @patch('src.investment.code_funcs.get_code_breakout')
    def test_reimbursement_level(self, mock_code_breakout, mock_methodology_breakout, 
                                mock_special_cases, mock_reimbursement_primary):
        # Setup
        mock_reimbursement_primary.return_value = [{"REIMB_TERM": "term1"}]
        mock_special_cases.return_value = [{"REIMB_TERM": "term1", "CARVEOUT_IND": "N"}]
        mock_methodology_breakout.return_value = [{"REIMB_TERM": "term1", "METHOD": "method1"}]
        mock_code_breakout.return_value = [{"REIMB_TERM": "term1", "CODE": "code1"}]
        
        mock_fields = MagicMock()
        mock_dataset = {}
        exhibit_page = "1"
        
        # Execute
        result = one_to_n_funcs.reimbursement_level(
            self.exhibit_text, self.filename, mock_fields, mock_dataset, exhibit_page
        )
        
        # Assert
        self.assertEqual(len(result), 1)
        self.assertEqual(result[0]["CODE"], "code1")
        mock_reimbursement_primary.assert_called_once()
        mock_special_cases.assert_called_once()
        mock_methodology_breakout.assert_called_once()
        mock_code_breakout.assert_called_once()

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