doczyai-pipelines/src/tests/test_postprocess.py

import unittest
from unittest.mock import MagicMock, patch

import pandas as pd
import pytest
from src.pipelines.shared.postprocessing.postprocessing_funcs import (
    deduplicate_provider_columns,
    flatten_singleton_string_list,
    format_rate_fields_with_commas,
    normalize_auto_renewal_term,
    normalize_cpt_fields,
    normalize_indicator_field,
    process_patient_age_range,
    remove_hyphens,
    remove_redundant_reimb_info,
    rename_columns,
    validate_and_reformat_date,
)


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):
        # Test numeric values
        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(1234.567), "1,234.57")
        self.assertEqual(format_rate_fields_with_commas(1000), "1,000.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("11"), "11")
        self.assertEqual(flatten_singleton_string_list("invalid"), "invalid")
        self.assertEqual(flatten_singleton_string_list(None), "")

    def test_rename_columns(self):
        """Tests the rename_columns function to ensure it correctly renames specified columns.

        Tests:
        1. Basic column renaming from PROCEDURE_CD to CPT4_PROC_CD
        2. Multiple columns being renamed
        3. Handling of columns not in the rename map
        4. Empty DataFrame
        """
        # Test case 1: Basic column renaming
        input_df1 = pd.DataFrame(
            {
                "PROCEDURE_CD": ["12345", "67890"],
                "PROCEDURE_CD_DESC": ["Test Procedure", "Another Procedure"],
                "OTHER_COLUMN": ["value1", "value2"],
            }
        )

        expected_df1 = pd.DataFrame(
            {
                "CPT4_PROC_CD": ["12345", "67890"],
                "CPT4_PROC_CD_DESC": ["Test Procedure", "Another Procedure"],
                "OTHER_COLUMN": ["value1", "value2"],
            }
        )

        result_df1 = rename_columns(input_df1)
        pd.testing.assert_frame_equal(result_df1, expected_df1)

        # Test case 2: Only some columns need renaming
        input_df2 = pd.DataFrame(
            {"PROCEDURE_CD": ["12345", "67890"], "OTHER_COLUMN": ["value1", "value2"]}
        )

        expected_df2 = pd.DataFrame(
            {"CPT4_PROC_CD": ["12345", "67890"], "OTHER_COLUMN": ["value1", "value2"]}
        )

        result_df2 = rename_columns(input_df2)
        pd.testing.assert_frame_equal(result_df2, expected_df2)

        # Test case 3: None of the columns need renaming
        input_df3 = pd.DataFrame({"COLUMN_A": ["a", "b"], "COLUMN_B": ["c", "d"]})

        result_df3 = rename_columns(input_df3)
        pd.testing.assert_frame_equal(result_df3, input_df3)  # Should be unchanged

        # Test case 4: Empty DataFrame
        empty_df = pd.DataFrame()
        result_empty_df = rename_columns(empty_df)
        pd.testing.assert_frame_equal(result_empty_df, empty_df)  # Should be unchanged

    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("['T0000-T9999, S0000-S9999']"),
            "['T0000-T9999', 'S0000-S9999']",
        )
        self.assertEqual(normalize_cpt_fields("123"), "['123']")
        self.assertEqual(normalize_cpt_fields(None), "")
        self.assertEqual(normalize_cpt_fields(123), "['123']")

    def test_process_patient_age_range(self):
        """Tests the process_patient_age_range function with various age range formats.

        Tests:
        1. Standard hyphenated ranges (e.g., "0-18")
        2. Single age values
        3. Text descriptions with "to"
        4. Text descriptions with "and under"
        5. Special cases like "newborn"
        6. Empty/None values
        7. Invalid formats
        8. Missing PATIENT_AGE_RANGE column
        """
        # Test case 1: DataFrame with PATIENT_AGE_RANGE column
        input_df = pd.DataFrame(
            {
                "PATIENT_AGE_RANGE": [
                    "0-18",  # Standard hyphenated range
                    "21",  # Single age
                ]
            }
        )

        result_df = process_patient_age_range(input_df)

        # Verify columns
        self.assertIn("PATIENT_AGE_MIN", result_df.columns)
        self.assertIn("PATIENT_AGE_MAX", result_df.columns)
        self.assertNotIn("PATIENT_AGE_RANGE", result_df.columns)

        # Expected values
        expected_min = ["0", "21"]
        expected_max = ["18", "21"]

        # Check transformations
        pd.testing.assert_series_equal(
            result_df["PATIENT_AGE_MIN"],
            pd.Series(expected_min, name="PATIENT_AGE_MIN"),
            check_dtype=False,
        )
        pd.testing.assert_series_equal(
            result_df["PATIENT_AGE_MAX"],
            pd.Series(expected_max, name="PATIENT_AGE_MAX"),
            check_dtype=False,
        )

        # Test case 2: DataFrame without PATIENT_AGE_RANGE column
        input_df_no_age = pd.DataFrame({"OTHER_COLUMN": ["value1", "value2"]})
        result_df_no_age = process_patient_age_range(input_df_no_age)

        # Verify the DataFrame is unchanged
        pd.testing.assert_frame_equal(input_df_no_age, result_df_no_age)

        # Test case 3: Empty DataFrame
        empty_df = pd.DataFrame()
        result_empty_df = process_patient_age_range(empty_df)

        # Verify empty DataFrame is unchanged
        pd.testing.assert_frame_equal(empty_df, result_empty_df)

    def test_validate_and_reformat_date(self):
        """Tests validate_and_reformat_date with various date formats.

        Tests:
        1. Date already in YYYY/MM/DD format
        2. Common alternative formats (YYYY-MM-DD, MM/DD/YYYY, etc.)
        3. Invalid date formats
        4. None and non-string values
        """
        # Date already in correct format
        self.assertEqual(validate_and_reformat_date("2023/01/15"), "2023/01/15")

        # Test various date formats that should be reformatted
        self.assertEqual(validate_and_reformat_date("2023-01-15"), "2023/01/15")
        self.assertEqual(validate_and_reformat_date("01/15/2023"), "2023/01/15")
        self.assertEqual(validate_and_reformat_date("15-Jan-2023"), "2023/01/15")

        # Test invalid formats - should return the original string
        self.assertEqual(validate_and_reformat_date("Invalid date"), "Invalid date")
        self.assertEqual(validate_and_reformat_date("01-15"), "01-15")

        # Test None and non-string values
        self.assertEqual(validate_and_reformat_date(None), None)
        self.assertEqual(validate_and_reformat_date(12345), 12345)

    def test_remove_redundant_reimb_info(self):
        """Tests remove_redundant_reimb_date function.

        Tests:
        1. When reimbursement dates match derived dates (should remove)
        2. When reimbursement dates differ from derived dates (should keep)
        3. When only some rows match (should remove only matching rows)
        4. When columns are missing (should return unchanged DataFrame)
        """
        # Test case 1: When dates match (should remove)
        input_df1 = pd.DataFrame(
            {
                "REIMB_EFFECTIVE_DT": ["2023/01/01", "2023/02/01"],
                "REIMB_TERMINATION_DT": ["2023/12/31", "2023/12/31"],
                "AARETE_DERIVED_EFFECTIVE_DT": ["2023/01/01", "2023/02/01"],
                "AARETE_DERIVED_TERMINATION_DT": ["2023/12/31", "2023/12/31"],
            }
        )

        expected_df1 = pd.DataFrame(
            {
                "REIMB_EFFECTIVE_DT": ["", ""],
                "REIMB_TERMINATION_DT": ["", ""],
                "AARETE_DERIVED_EFFECTIVE_DT": ["2023/01/01", "2023/02/01"],
                "AARETE_DERIVED_TERMINATION_DT": ["2023/12/31", "2023/12/31"],
            }
        )

        result_df1 = remove_redundant_reimb_info(input_df1)
        pd.testing.assert_frame_equal(result_df1, expected_df1)

        # Test case 2: When dates differ (should keep)
        input_df2 = pd.DataFrame(
            {
                "REIMB_EFFECTIVE_DT": ["2023/01/15", "2023/02/15"],
                "REIMB_TERMINATION_DT": ["2023/12/15", "2023/12/15"],
                "AARETE_DERIVED_EFFECTIVE_DT": ["2023/01/01", "2023/02/01"],
                "AARETE_DERIVED_TERMINATION_DT": ["2023/12/31", "2023/12/31"],
            }
        )

        # Result should be unchanged
        result_df2 = remove_redundant_reimb_info(input_df2)
        pd.testing.assert_frame_equal(result_df2, input_df2)

        # Test case 3: Mixed case - some match, some don't
        input_df3 = pd.DataFrame(
            {
                "REIMB_EFFECTIVE_DT": ["2023/01/01", "2023/02/15"],
                "REIMB_TERMINATION_DT": ["2023/12/31", "2023/12/15"],
                "AARETE_DERIVED_EFFECTIVE_DT": ["2023/01/01", "2023/02/01"],
                "AARETE_DERIVED_TERMINATION_DT": ["2023/12/31", "2023/12/31"],
            }
        )

        expected_df3 = pd.DataFrame(
            {
                "REIMB_EFFECTIVE_DT": ["", "2023/02/15"],
                "REIMB_TERMINATION_DT": ["", "2023/12/15"],
                "AARETE_DERIVED_EFFECTIVE_DT": ["2023/01/01", "2023/02/01"],
                "AARETE_DERIVED_TERMINATION_DT": ["2023/12/31", "2023/12/31"],
            }
        )

        result_df3 = remove_redundant_reimb_info(input_df3)
        pd.testing.assert_frame_equal(result_df3, expected_df3)

        # Test case 4: Missing columns
        input_df4 = pd.DataFrame(
            {
                "REIMB_EFFECTIVE_DT": ["2023/01/01", "2023/02/01"],
                "OTHER_COLUMN": ["value1", "value2"],
            }
        )

        # Result should be unchanged
        result_df4 = remove_redundant_reimb_info(input_df4)
        pd.testing.assert_frame_equal(result_df4, input_df4)

    def test_deduplicate_provider_columns(self):
        """Tests deduplicate_provider_columns function.

        Tests:
        1. Basic deduplication - removes GROUP values from OTHER fields
        2. Internal deduplication - removes duplicates within OTHER fields
        3. Combined scenario - both GROUP removal and internal deduplication
        4. Empty/missing values handling
        5. Missing columns - should return unchanged DataFrame
        6. Empty DataFrame
        """

        # Test case 1: Basic GROUP removal
        # deduplicate_provider_columns now works with lists, not pipe-delimited strings
        input_df1 = pd.DataFrame(
            {
                "PROV_GROUP_TIN": [["123456789"]],
                "PROV_GROUP_NPI": [["1234567890"]],
                "PROV_GROUP_NAME_FULL": [["Main Hospital"]],
                "PROV_OTHER_TIN": [["123456789", "987654321"]],
                "PROV_OTHER_NPI": [["1234567890", "0987654321"]],
                "PROV_OTHER_NAME_FULL": [["Main Hospital", "Other Clinic"]],
            }
        )

        expected_df1 = pd.DataFrame(
            {
                "PROV_GROUP_TIN": [["123456789"]],
                "PROV_GROUP_NPI": [["1234567890"]],
                "PROV_GROUP_NAME_FULL": [["Main Hospital"]],
                "PROV_OTHER_TIN": [["987654321"]],
                "PROV_OTHER_NPI": [["0987654321"]],
                "PROV_OTHER_NAME_FULL": [["Other Clinic"]],
            }
        )

        result_df1 = deduplicate_provider_columns(input_df1)
        pd.testing.assert_frame_equal(result_df1, expected_df1)

        # Test case 2: Internal deduplication (your original example)
        # deduplicate_provider_columns now works with lists, not pipe-delimited strings
        input_df2 = pd.DataFrame(
            {
                "PROV_GROUP_TIN": [["061798267"]],
                "PROV_GROUP_NPI": [["1111111111"]],
                "PROV_GROUP_NAME_FULL": [["Group Practice"]],
                "PROV_OTHER_TIN": [
                    [
                        "061798267",
                        "061798267",
                        "UNKNOWN",
                        "UNKNOWN",
                        "061992277",
                        "061798267",
                    ]
                ],
                "PROV_OTHER_NPI": [
                    ["1111111111", "2222222222", "2222222222", "UNKNOWN"]
                ],
                "PROV_OTHER_NAME_FULL": [
                    ["Group Practice", "Other Practice", "Other Practice", "UNKNOWN"]
                ],
            }
        )

        expected_df2 = pd.DataFrame(
            {
                "PROV_GROUP_TIN": [["061798267"]],
                "PROV_GROUP_NPI": [["1111111111"]],
                "PROV_GROUP_NAME_FULL": [["Group Practice"]],
                "PROV_OTHER_TIN": [["061992277"]],
                "PROV_OTHER_NPI": [["2222222222"]],
                "PROV_OTHER_NAME_FULL": [["Other Practice"]],
            }
        )

        result_df2 = deduplicate_provider_columns(input_df2)
        pd.testing.assert_frame_equal(result_df2, expected_df2)

        # Test case 3: Empty OTHER fields after deduplication
        # deduplicate_provider_columns now works with lists, not pipe-delimited strings
        input_df3 = pd.DataFrame(
            {
                "PROV_GROUP_TIN": [["123456789"]],
                "PROV_GROUP_NPI": [["1234567890"]],
                "PROV_GROUP_NAME_FULL": [["Main Hospital"]],
                "PROV_OTHER_TIN": [["123456789", "123456789", "UNKNOWN"]],
                "PROV_OTHER_NPI": [["1234567890", "UNKNOWN", "UNKNOWN"]],
                "PROV_OTHER_NAME_FULL": [["Main Hospital", "UNKNOWN"]],
            }
        )

        expected_df3 = pd.DataFrame(
            {
                "PROV_GROUP_TIN": [["123456789"]],
                "PROV_GROUP_NPI": [["1234567890"]],
                "PROV_GROUP_NAME_FULL": [["Main Hospital"]],
                "PROV_OTHER_TIN": [[]],
                "PROV_OTHER_NPI": [[]],
                "PROV_OTHER_NAME_FULL": [[]],
            }
        )

        result_df3 = deduplicate_provider_columns(input_df3)
        pd.testing.assert_frame_equal(result_df3, expected_df3)

        # Test case 4: Multiple rows
        # deduplicate_provider_columns now works with lists, not pipe-delimited strings
        input_df4 = pd.DataFrame(
            {
                "PROV_GROUP_TIN": [["111111111"], ["222222222"]],
                "PROV_GROUP_NPI": [["1111111111"], ["2222222222"]],
                "PROV_GROUP_NAME_FULL": [["Hospital A"], ["Hospital B"]],
                "PROV_OTHER_TIN": [
                    ["111111111", "333333333"],
                    ["444444444", "222222222", "444444444"],
                ],
                "PROV_OTHER_NPI": [
                    ["3333333333", "1111111111"],
                    ["4444444444", "2222222222"],
                ],
                "PROV_OTHER_NAME_FULL": [
                    ["Clinic C", "Hospital A"],
                    ["Clinic D", "Hospital B"],
                ],
            }
        )

        expected_df4 = pd.DataFrame(
            {
                "PROV_GROUP_TIN": [["111111111"], ["222222222"]],
                "PROV_GROUP_NPI": [["1111111111"], ["2222222222"]],
                "PROV_GROUP_NAME_FULL": [["Hospital A"], ["Hospital B"]],
                "PROV_OTHER_TIN": [["333333333"], ["444444444"]],
                "PROV_OTHER_NPI": [["3333333333"], ["4444444444"]],
                "PROV_OTHER_NAME_FULL": [["Clinic C"], ["Clinic D"]],
            }
        )

        result_df4 = deduplicate_provider_columns(input_df4)
        pd.testing.assert_frame_equal(result_df4, expected_df4)

        # Test case 5: Missing columns - should return unchanged
        input_df5 = pd.DataFrame(
            {"PROV_GROUP_TIN": ["123456789"], "OTHER_COLUMN": ["value1"]}
        )

        result_df5 = deduplicate_provider_columns(input_df5)
        pd.testing.assert_frame_equal(result_df5, input_df5)

        # Test case 6: Empty DataFrame
        empty_df = pd.DataFrame()
        result_empty_df = deduplicate_provider_columns(empty_df)
        pd.testing.assert_frame_equal(empty_df, result_empty_df)

        # Test case 7: Empty OTHER fields (already empty lists)
        # deduplicate_provider_columns now works with lists, not pipe-delimited strings
        input_df7 = pd.DataFrame(
            {
                "PROV_GROUP_TIN": [["123456789"]],
                "PROV_GROUP_NPI": [["1234567890"]],
                "PROV_GROUP_NAME_FULL": [["Main Hospital"]],
                "PROV_OTHER_TIN": [[]],
                "PROV_OTHER_NPI": [[]],
                "PROV_OTHER_NAME_FULL": [[]],
            }
        )

        result_df7 = deduplicate_provider_columns(input_df7)
        pd.testing.assert_frame_equal(result_df7, input_df7)


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