doczyai-pipelines/fieldExtraction/tests/test_dynamic.py

import json
import os
import unittest
from unittest.mock import MagicMock, patch

import pandas as pd
from constants.constants import Constants
from src.investment import dynamic_funcs, prompt_calls
from src.prompts.fieldset import Field, FieldSet


class TestDynamicFuncsWithRealConstants(unittest.TestCase):
    """Unit tests for functions in src/investment/dynamic_funcs.py using real Constants."""

    def setUp(self):
        """Set up test fixtures before each test method."""
        # Create real Constants object
        self.constants = Constants()

        # Create test fields
        self.test_field = Field.from_values(
            field_name="TEST_FIELD",
            relationship="one-to-one",
            field_type="exhibit_level",
            prompt="What is the test field?",
            base_field="TEST_FIELD_BASE",
        )

        self.test_field2 = Field.from_values(
            field_name="TEST_FIELD2",
            relationship="one-to-many",
            field_type="exhibit_level",
            prompt="What is the second test field?",
            base_field="TEST_FIELD2_BASE",
        )

        self.escalator_field = Field.from_values(
            field_name="RATE_ESCALATOR_TERM",
            relationship="one-to-many",
            field_type="rate_escalator",
            prompt="What is the rate escalator?",
            base_field="RATE_ESCALATOR_STATEMENT",
        )

        self.outlier_field = Field.from_values(
            field_name="OUTLIER_TERMS",
            relationship="one-to-many",
            field_type="exhibit_level",
            prompt="What are the outlier terms?",
            base_field="OUTLIER_TERMS",
        )

        # Create test FieldSets
        self.test_fieldset = FieldSet()
        self.test_fieldset.add_field(self.test_field)
        self.test_fieldset.add_field(self.test_field2)

        # Mock filename
        self.filename = "test_file.pdf"

        # Sample text
        self.exhibit_text = "This is a sample exhibit text with test data."
        self.exhibit_header = "This is a sample exhibit header with test data."

        # Sample answer dictionaries
        self.answer_dicts = [
            {
                "TEST_FIELD": "Value 1",
                "TEST_FIELD2": "",
                "RATE_ESCALATOR_TERM": "",
                "AARETE_DERIVED_LOB": "Commercial",
            },
            {
                "TEST_FIELD": "Value 2",
                "TEST_FIELD2": "",
                "RATE_ESCALATOR_TERM": "",
                "AARETE_DERIVED_LOB": "Medicare",
            },
            {
                "TEST_FIELD": "Value 3",
                "TEST_FIELD2": "",
                "OUTLIER_TERMS": "",
                "AARETE_DERIVED_LOB": "",
            },
        ]

    @patch("src.utils.llm_utils.invoke_claude")
    @patch("src.utils.string_utils.universal_json_load")
    def test_prompt_dynamic(self, mock_json_load, mock_invoke_claude):
        """Test the prompt_dynamic function."""
        # Setup
        field_prompts = {
            "TEST_FIELD": "What is the test field?",
            "TEST_FIELD2": "What is the second test field?",
        }

        # Configure mocks
        mock_invoke_claude.return_value = (
            '{"TEST_FIELD": ["Value 1"], "TEST_FIELD2": ["Value 2"]}'
        )
        mock_json_load.return_value = {
            "TEST_FIELD": ["Value 1"],
            "TEST_FIELD2": ["Value 2"],
        }

        # Execute
        result = prompt_calls.prompt_dynamic(
            self.exhibit_text, field_prompts, self.filename
        )

        # Assert
        mock_invoke_claude.assert_called_once()
        mock_json_load.assert_called_once_with(
            '{"TEST_FIELD": ["Value 1"], "TEST_FIELD2": ["Value 2"]}'
        )
        self.assertEqual(
            result, {"TEST_FIELD": ["Value 1"], "TEST_FIELD2": ["Value 2"]}
        )

        # Verify that the prompt contains our text
        prompt_call = mock_invoke_claude.call_args[0][0]
        self.assertIn(self.exhibit_text, prompt_call)
        self.assertIn("TEST_FIELD", prompt_call)
        self.assertIn("TEST_FIELD2", prompt_call)

    def test_add_full_context_field_normal(self):
        """Test adding a normal field to one_to_one_fields."""
        # Setup
        one_to_one_fields = FieldSet()
        field_to_add = self.test_field

        # Execute
        result = dynamic_funcs.add_one_to_one_field(
            one_to_one_fields, field_to_add, self.answer_dicts, self.constants
        )

        # Assert
        self.assertEqual(len(result.fields), 1)
        self.assertEqual(result.fields[0].field_name, "TEST_FIELD")
        self.assertEqual(result.fields[0].field_type, "smart_chunked")
        self.assertEqual(result.fields[0].relationship, "one_to_one")
        
    def test_add_full_context_field_program_with_single_lob(self):
        """Test that program fields are added when only one LOB exists."""
        # Setup
        one_to_one_fields = FieldSet()
        program_field = Field.from_values(
            field_name="PROGRAM",
            relationship="one-to-one",
            field_type="exhibit_level",
            prompt="What is the program?",
            base_field="PROGRAM",
        )

        # Use modified answer_dicts with only one non-empty LOB
        single_lob_answer_dicts = [
            {
                "TEST_FIELD": "Value 1",
                "TEST_FIELD2": "",
                "RATE_ESCALATOR_TERM": "",
                "AARETE_DERIVED_LOB": "Commercial",
            },
            {
                "TEST_FIELD": "Value 2",
                "TEST_FIELD2": "",
                "RATE_ESCALATOR_TERM": "",
                "AARETE_DERIVED_LOB": "",
            },
            {
                "TEST_FIELD": "Value 3",
                "TEST_FIELD2": "",
                "OUTLIER_TERMS": "",
                "AARETE_DERIVED_LOB": "",
            },
        ]

        # Execute
        result = dynamic_funcs.add_one_to_one_field(
            one_to_one_fields, program_field, single_lob_answer_dicts, self.constants
        )

        # Assert
        self.assertEqual(len(result.fields), 1)  # Field should be added
        self.assertEqual(result.fields[0].field_name, "PROGRAM")

    @patch("src.investment.dynamic_funcs.Field.load_from_file")
    @patch("src.investment.dynamic_funcs.FieldSet")
    def test_get_dynamic_one_to_one_fields_any_empty(
        self, mock_fieldset_class, mock_load_field
    ):
        """Test getting dynamic one-to-one fields with any empty criterion."""
        # Setup
        # Configure mock for fieldset class
        mock_all_empty_fields = MagicMock()
        mock_all_empty_fields.fields = []

        mock_any_empty_fields = MagicMock()
        mock_field1 = MagicMock()
        mock_field1.field_name = "RATE_ESCALATOR_TERM"
        mock_field1.base_field = "RATE_ESCALATOR_STATEMENT"
        mock_any_empty_fields.fields = [mock_field1]

        mock_one_to_one_fields = MagicMock()
        mock_one_to_one_fields.fields = []

        mock_fieldset_class.side_effect = [
            mock_all_empty_fields,
            mock_any_empty_fields,
            mock_one_to_one_fields,
            MagicMock(), # extra FieldSet() call safeguard
            MagicMock(), # optional safety if code evolves
        ]

        # Configure mock for loading fields
        escalator_statement_field = Field.from_values(
            field_name="RATE_ESCALATOR_STATEMENT",
            relationship="one-to-one",
            field_type="full_context",
            prompt="What is the rate escalator statement?",
        )
        mock_load_field.return_value = escalator_statement_field

        # Execute
        result = dynamic_funcs.get_dynamic_one_to_one_fields(self.answer_dicts, self.constants)
        
    @patch("src.investment.prompt_calls.prompt_dynamic")
    def test_dynamic_with_real_constants(self, mock_prompt_dynamic):
        """Test dynamic function with real Constants object."""
        # Setup
        dynamic_fields = self.test_fieldset
        reimbursement_level_fields = FieldSet()
        exhibit_level_answers = {}

        # Configure mocks
        mock_prompt_dynamic.side_effect = [
            {"TEST_FIELD": "Header Answer", "TEST_FIELD2": ""},  # Header results
            {"TEST_FIELD2": "Text Answer"},  # Text results (only for remaining field)
        ]

        # Execute
        exhibit_level_answers, updated_reimbursement_fields = dynamic_funcs.dynamic(
            self.exhibit_text,
            self.exhibit_header,
            exhibit_level_answers,
            self.constants,
            self.filename,
            dynamic_fields,
            reimbursement_level_fields,
        )

        # Assert
        self.assertEqual(
            len(updated_reimbursement_fields.fields), 2
        )  # Both fields should be added to reimbursement level

        # Check that the first field was processed from header results
        self.assertEqual(
            updated_reimbursement_fields.fields[0].field_name, "TEST_FIELD"
        )
        self.assertEqual(
            updated_reimbursement_fields.fields[0].valid_values, "Header Answer"
        )

        # Check that the second field was processed from text results
        self.assertEqual(
            updated_reimbursement_fields.fields[1].field_name, "TEST_FIELD2"
        )
        self.assertEqual(
            updated_reimbursement_fields.fields[1].valid_values, "Text Answer"
        )

        # Verify that prompt_dynamic was called twice
        self.assertEqual(mock_prompt_dynamic.call_count, 2)

    @patch("src.investment.prompt_calls.prompt_dynamic")
    def test_dynamic_with_no_fields(self, mock_prompt_dynamic):
        """Test dynamic function with an empty fieldset."""
        # Setup
        empty_fieldset = FieldSet()
        reimbursement_level_fields = FieldSet()
        exhibit_level_answers = {}

        # Execute
        exhibit_level_answers, updated_reimbursement_fields = dynamic_funcs.dynamic(
            self.exhibit_text,
            self.exhibit_header,
            exhibit_level_answers,
            self.constants,
            self.filename,
            empty_fieldset,
            reimbursement_level_fields,
        )

        # Assert
        self.assertEqual(len(exhibit_level_answers), 0)
        self.assertEqual(len(updated_reimbursement_fields.fields), 0)
        mock_prompt_dynamic.assert_not_called()

    @patch("src.investment.prompt_calls.prompt_dynamic")
    def test_dynamic_with_no_answers(self, mock_prompt_dynamic):
        """Test dynamic function when neither header nor text contains answers."""
        # Setup
        dynamic_fields = self.test_fieldset
        reimbursement_level_fields = FieldSet()
        exhibit_level_answers = {}

        # Configure mocks - no answers anywhere
        mock_prompt_dynamic.side_effect = [
            {"TEST_FIELD": "", "TEST_FIELD2": ""},  # Header results
            {"TEST_FIELD": "", "TEST_FIELD2": ""},  # Text results
        ]

        # Execute
        exhibit_level_answers, updated_reimbursement_fields = dynamic_funcs.dynamic(
            self.exhibit_text,
            self.exhibit_header,
            exhibit_level_answers,
            self.constants,
            self.filename,
            dynamic_fields,
            reimbursement_level_fields,
        )

        # Assert
        self.assertEqual(
            len(updated_reimbursement_fields.fields), 0
        )  # No fields should be added to reimbursement level
        self.assertEqual(
            exhibit_level_answers, {"TEST_FIELD": "", "TEST_FIELD2": ""}
        )  # Both fields should be added to exhibit level
        self.assertEqual(mock_prompt_dynamic.call_count, 2)


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