doczyai-pipelines/fieldExtraction/tests/investment_postprocess_test.py

from unittest.mock import patch, MagicMock
import unittest
from unittest.mock import patch
import pytest
import pandas as pd
from src.investment.investment_postprocessing_funcs import (
    normalize_indicator_field,
    format_rate_fields_with_commas,
    remove_hyphens,
    flatten_singleton_string_list,
    validate_and_reformat_date,
    date_postprocess,
    normalize_auto_renewal_term,
    normalize_cpt_fields
)


class TestPostprocessFunctions(unittest.TestCase):

    def test_normalize_indicator_field(self):
        self.assertEqual(normalize_indicator_field("Y"), "Y")
        self.assertEqual(normalize_indicator_field("y"), "Y")
        self.assertEqual(normalize_indicator_field("N"), "N")
        self.assertEqual(normalize_indicator_field(""), "N")
        self.assertEqual(normalize_indicator_field(None), "N")

    def test_format_rate_fields_with_commas(self):
        self.assertEqual(format_rate_fields_with_commas("1234.567"), "1,234.57")
        self.assertEqual(format_rate_fields_with_commas("1000"), "1,000.00")
        self.assertEqual(format_rate_fields_with_commas("abc"), "0.00")
        self.assertEqual(format_rate_fields_with_commas(None), "0.00")

    def test_remove_hyphens(self):
        self.assertEqual(remove_hyphens("123-45-6789"), "123456789")
        self.assertEqual(remove_hyphens("123-456"), "123456")
        self.assertEqual(remove_hyphens(None), "")
        self.assertEqual(remove_hyphens(""), "")

    def test_flatten_singleton_string_list(self):
        self.assertEqual(flatten_singleton_string_list("['123']"), "123")
        self.assertEqual(flatten_singleton_string_list("['123', '456']"), ['123', '456'])
        self.assertEqual(flatten_singleton_string_list("invalid"), "invalid")
        self.assertEqual(flatten_singleton_string_list(None), None)


    @patch("src.utils.llm_utils.invoke_claude")
    @patch("src.investment.investment_postprocessing_funcs.invoke_derived_term_date")
    @patch("src.prompts.investment_prompts.FieldSet.load_from_file")  # Patch this specific method
    def test_date_postprocess(self, mock_load_from_file, mock_derived_term_date, mock_invoke_claude):
        """Tests the date_postprocess function with mocked dependencies.
        
        This test verifies that:
        1. Date fields are standardized to YYYYMMDD format using the LLM (mocked)
        2. The derived termination date is properly calculated and added to the DataFrame
        
        The test uses multiple patches to isolate the function from external dependencies:
        - Mocks `invoke_claude` to return a fixed date string ("|20230101|")  
        - Mocks `invoke_derived_term_date` to return a fixed termination date ("20231231")
        - Mocks `FieldSet.load_from_file` and `FieldSet.filter` to avoid file operations
          and provide controlled field definitions
    
        The test checks that the date fields in the DataFrame are correctly formatted and
        that the derived termination date is added as expected.
        """
        # Setup field mock
        field_mocks = [
            MagicMock(field_name="CONTRACT_EFFECTIVE_DT"),
            MagicMock(field_name="CONTRACT_TERMINATION_DT"),
            MagicMock(field_name="TERMINATION_DT")
        ]
        
        # Configure the mock methods
        mock_invoke_claude.return_value = "|20230101|"
        mock_derived_term_date.return_value = "20231231"
        
        df = pd.DataFrame({
            "CONTRACT_EFFECTIVE_DT": ["2023-01-01", "01/01/2023", "invalid"],
            "CONTRACT_TERMINATION_DT": ["2023-12-31", "12/31/2023", "invalid"],
            "TERMINATION_DT": ["2023-12-31", "12/31/2023", "invalid"],
            "AARETE_DERIVED_EFFECTIVE_DT": ["2023-01-01", "2023-01-01", "2023-01-01"]
        })

        with patch("src.prompts.investment_prompts.FieldSet.filter") as mock_filter:
            # Configure the filter to return a mock object with fields attribute
            mock_filter.return_value = MagicMock(fields=field_mocks)
            
            result_df = date_postprocess(df, "mock_path")

        assert result_df["CONTRACT_EFFECTIVE_DT"].tolist() == ["20230101", "20230101", "20230101"]
        assert result_df["AARETE_DERIVED_TERMINATION_DT"].tolist() == ["20231231", "20231231", "20231231"]

    def test_normalize_auto_renewal_term(self):
        self.assertEqual(normalize_auto_renewal_term("12 months"), "1 year")
        self.assertEqual(normalize_auto_renewal_term("one year"), "1 year")
        self.assertEqual(normalize_auto_renewal_term("month to month"), "1 month")
        self.assertEqual(normalize_auto_renewal_term("(12) 12 months"), "1 year")
        self.assertEqual(normalize_auto_renewal_term(None), "")

    def test_normalize_cpt_fields(self):
        self.assertEqual(normalize_cpt_fields("[123, 456]"), ["123", "456"])
        self.assertEqual(normalize_cpt_fields("123-456"), ["123-456"])
        self.assertEqual(normalize_cpt_fields("123"), ["123"])
        self.assertEqual(normalize_cpt_fields(None), [])
        self.assertEqual(normalize_cpt_fields(123), ["123"])


@pytest.mark.parametrize(
    "input_date, expected_output",
    [
        # Valid YYYY/MM/DD format
        ("2023/10/15", "2023/10/15"),
        # Valid formats to be reformatted to YYYY/MM/DD
        ("2023-10-15", "2023/10/15"),
        ("10/15/2023", "2023/10/15"),
        ("15-Oct-2023", "2023/10/15"),
        ("15/10/2023", "2023/10/15"),
        # Invalid formats or non-date strings
        ("15-10-2023", "15-10-2023"),
        ("InvalidDate", "InvalidDate"),
        # Non-string input
        (None, None),
        (12345, 12345),
        # Edge cases
        ("", ""),
        ("2023/13/01", "2023/13/01"),  # Invalid month
        ("2023/02/30", "2023/02/30"),  # Invalid day
    ],
)
def test_validate_and_reformat_date(input_date, expected_output):
    assert validate_and_reformat_date(input_date) == expected_output

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