doczyai-pipelines/fieldExtraction/tests/test_one_to_n.py

import unittest
from unittest.mock import Mock, patch, MagicMock
from src.investment import one_to_n_funcs
from src.prompts.fieldset import FieldSet
from constants.constants import Constants


class TestOneToNFuncs(unittest.TestCase):
    """Test suite for one_to_n_funcs.py functions."""

    def setUp(self):
        """Set up test fixtures."""
        self.constants = Constants()
        self.filename = "test_contract.pdf"
        self.exhibit_text = "This is a sample exhibit text with reimbursement terms."
        self.exhibit_header = "Exhibit A - Fee Schedule"
        self.exhibit_page = "1"

    # ==================== exhibit_level Tests ====================

    @patch('src.utils.llm_utils.invoke_claude')  # Mock at the invoke_claude level
    @patch('src.investment.one_to_n_funcs.dynamic_funcs.dynamic_primary')
    def test_exhibit_level_basic(self, mock_dynamic_primary, mock_invoke_claude):
        """Test basic exhibit level processing."""
        # Mock LLM response for dynamic_primary call
        mock_invoke_claude.return_value = '{"AARETE_DERIVED_LOB": "Medicaid", "AARETE_DERIVED_NETWORK": "In-Network"}'
        
        # Mock dynamic primary returning updated answers and fields
        mock_dynamic_fields = Mock(spec=FieldSet)
        mock_dynamic_primary.return_value = (
            {"SERVICE_CATEGORY_CD": "INPATIENT", "PROV_TYPE_CD": "HOSPITAL", "AARETE_DERIVED_LOB": "Medicaid"},
            mock_dynamic_fields
        )

        exhibit_level_answers, dynamic_fields = one_to_n_funcs.exhibit_level(
            self.exhibit_text, self.exhibit_header, self.exhibit_page, self.constants, self.filename
        )
        
        # Assertions
        assert exhibit_level_answers is not None
        mock_invoke_claude.assert_called()

    # ==================== reimbursement_level Tests ====================

    @patch('src.investment.one_to_n_funcs.prompt_calls.prompt_reimbursement_primary')
    @patch('src.investment.one_to_n_funcs.clean_reimbursement_primary')
    def test_reimbursement_level_with_results(self, mock_clean, mock_prompt_reimb):
        """Test reimbursement level extraction with valid results."""
        mock_prompt_reimb.return_value = [
            {"SERVICE_TERM": "Lab Services", "REIMB_TERM": "100% of Medicare"},
            {"SERVICE_TERM": "Radiology", "REIMB_TERM": "90% of billed charges"}
        ]
        mock_clean.return_value = mock_prompt_reimb.return_value

        page_text = "Lab Services will be reimbursed at 100% of Medicare."
        result = one_to_n_funcs.reimbursement_level(page_text, self.constants, self.filename)

        self.assertEqual(len(result), 2)
        mock_prompt_reimb.assert_called_once_with(page_text, self.filename)
        mock_clean.assert_called_once()

    @patch('src.investment.one_to_n_funcs.prompt_calls.prompt_reimbursement_primary')
    def test_reimbursement_level_no_results(self, mock_prompt_reimb):
        """Test reimbursement level when no reimbursements found."""
        mock_prompt_reimb.return_value = None

        page_text = "This page has no reimbursement information."
        result = one_to_n_funcs.reimbursement_level(page_text, self.constants, self.filename)

        self.assertIsNone(result)

    # ==================== clean_reimbursement_primary Tests ====================

    @patch('src.investment.one_to_n_funcs.filter_services_without_reimbursements')
    def test_clean_reimbursement_primary_basic(self, mock_filter):
        """Test basic cleaning of reimbursement primary answers."""
        input_answers = [
            {"SERVICE_TERM": "Labs", "REIMB_TERM": "100% of Medicare"},
            {"SERVICE_TERM": "X-Ray", "REIMB_TERM": "Fee schedule"}
        ]
        mock_filter.return_value = input_answers

        result = one_to_n_funcs.clean_reimbursement_primary(
            input_answers, self.constants, self.filename
        )

        self.assertEqual(len(result), 2)
        mock_filter.assert_called_once_with(input_answers, self.filename)

    def test_clean_reimbursement_primary_empty_input(self):
        """Test cleaning with empty input."""
        result = one_to_n_funcs.clean_reimbursement_primary(
            [], self.constants, self.filename
        )
        self.assertEqual(result, [])

    def test_clean_reimbursement_primary_none_input(self):
        """Test cleaning with None input."""
        result = one_to_n_funcs.clean_reimbursement_primary(
            None, self.constants, self.filename
        )
        self.assertEqual(result, [])

    # ==================== methodology_breakout Tests ====================

    @patch('src.investment.one_to_n_funcs.methodology_breakout_single_row')
    def test_methodology_breakout_multiple_rows(self, mock_single_row):
        """Test methodology breakout with multiple rows."""
        mock_single_row.side_effect = [
            [{"SERVICE_TERM": "Labs", "AARETE_DERIVED_REIMB_METHOD": "Fee Schedule"}],
            [{"SERVICE_TERM": "Surgery", "AARETE_DERIVED_REIMB_METHOD": "Grouper"}]
        ]

        input_answers = [
            {"SERVICE_TERM": "Labs", "REIMB_TERM": "Medicare fee schedule"},
            {"SERVICE_TERM": "Surgery", "REIMB_TERM": "MS-DRG"}
        ]

        result = one_to_n_funcs.methodology_breakout(
            input_answers, self.constants, self.filename
        )

        self.assertEqual(len(result), 2)
        self.assertEqual(mock_single_row.call_count, 2)

    # ==================== methodology_breakout_single_row Tests ====================

    @patch('src.investment.one_to_n_funcs.prompt_calls.prompt_methodology_breakout')
    @patch('src.investment.one_to_n_funcs.methodology_breakout_secondary')
    def test_methodology_breakout_single_row_basic(self, mock_secondary, mock_prompt):
        """Test single row methodology breakout."""
        mock_prompt.return_value = {
            "AARETE_DERIVED_REIMB_METHOD": "Fee Schedule",
            "REIMB_FEE_RATE": "N/A",
            "REIMB_PCT_RATE": "100"
        }
        mock_secondary.return_value = {"FEE_SCHEDULE": "Medicare"}

        answer_dict = {"SERVICE_TERM": "Labs", "REIMB_TERM": "100% of Medicare"}
        
        result = one_to_n_funcs.methodology_breakout_single_row(
            answer_dict, self.constants, self.filename
        )

        self.assertEqual(len(result), 1)
        self.assertIn("AARETE_DERIVED_REIMB_METHOD", result[0])
        self.assertIn("FEE_SCHEDULE", result[0])

    @patch('src.investment.one_to_n_funcs.prompt_calls.prompt_methodology_breakout')
    def test_methodology_breakout_single_row_no_reimb_term(self, mock_prompt):
        """Test single row breakout when REIMB_TERM is missing."""
        answer_dict = {"SERVICE_TERM": "Labs"}
        
        result = one_to_n_funcs.methodology_breakout_single_row(
            answer_dict, self.constants, self.filename
        )

        self.assertEqual(result, [answer_dict])
        mock_prompt.assert_not_called()

    @patch('src.investment.one_to_n_funcs.prompt_calls.prompt_methodology_breakout')
    @patch('src.investment.one_to_n_funcs.methodology_breakout_secondary')
    def test_methodology_breakout_single_row_list_output(self, mock_secondary, mock_prompt):
        """Test single row breakout when LLM returns list of methodologies."""
        mock_prompt.return_value = [
            {"AARETE_DERIVED_REIMB_METHOD": "Fee Schedule"},
            {"AARETE_DERIVED_REIMB_METHOD": "Per Diem"}
        ]
        mock_secondary.side_effect = [
            {"FEE_SCHEDULE": "Medicare"},
            {"FEE_SCHEDULE": "N/A"}
        ]

        answer_dict = {"SERVICE_TERM": "Labs", "REIMB_TERM": "Medicare or $500 per day"}
        
        result = one_to_n_funcs.methodology_breakout_single_row(
            answer_dict, self.constants, self.filename
        )

        self.assertEqual(len(result), 2)
        self.assertEqual(mock_secondary.call_count, 2)

    # ==================== methodology_breakout_secondary Tests ====================

    @patch('src.investment.one_to_n_funcs.prompt_calls.prompt_fee_schedule_breakout')
    def test_methodology_breakout_secondary_fee_schedule(self, mock_fs_breakout):
        """Test secondary breakout for Fee Schedule methodology."""
        mock_fs_breakout.return_value = {
            "FEE_SCHEDULE": "Medicare",
            "FEE_SCHEDULE_VERSION": "2024"
        }

        methodology_dict = {"AARETE_DERIVED_REIMB_METHOD": "Fee Schedule"}
        
        with patch('src.investment.one_to_n_funcs.FieldSet') as mock_fieldset:
            mock_grouper_fields = Mock()
            mock_grouper_fields.list_fields.return_value = ["GROUPER_TYPE", "DRG_WEIGHT_TYPE"]
            mock_fieldset.return_value = mock_grouper_fields

            result = one_to_n_funcs.methodology_breakout_secondary(
                methodology_dict, "Labs", "Medicare fee schedule",
                Mock(), mock_grouper_fields, self.constants, self.filename
            )

        self.assertIn("FEE_SCHEDULE", result)
        self.assertEqual(result["GROUPER_TYPE"], "N/A")
        self.assertEqual(result["DRG_WEIGHT_TYPE"], "N/A")

    @patch('src.investment.one_to_n_funcs.prompt_calls.prompt_grouper_breakout')
    def test_methodology_breakout_secondary_grouper(self, mock_grouper_breakout):
        """Test secondary breakout for Grouper methodology."""
        mock_grouper_breakout.return_value = {
            "GROUPER_TYPE": "MS-DRG",
            "DRG_WEIGHT_TYPE": "CMS"
        }

        methodology_dict = {"AARETE_DERIVED_REIMB_METHOD": "Grouper"}
        
        with patch('src.investment.one_to_n_funcs.FieldSet') as mock_fieldset:
            mock_fs_fields = Mock()
            mock_fs_fields.list_fields.return_value = ["FEE_SCHEDULE", "FEE_SCHEDULE_VERSION"]
            mock_fieldset.return_value = mock_fs_fields

            result = one_to_n_funcs.methodology_breakout_secondary(
                methodology_dict, "Surgery", "MS-DRG",
                mock_fs_fields, Mock(), self.constants, self.filename
            )

        self.assertIn("GROUPER_TYPE", result)
        self.assertEqual(result["FEE_SCHEDULE"], "N/A")
        self.assertEqual(result["FEE_SCHEDULE_VERSION"], "N/A")

    def test_methodology_breakout_secondary_other_methodology(self):
        """Test secondary breakout for non-FS, non-Grouper methodology."""
        methodology_dict = {"AARETE_DERIVED_REIMB_METHOD": "Per Diem"}
        
        with patch('src.investment.one_to_n_funcs.FieldSet') as mock_fieldset:
            mock_fields = Mock()
            mock_fields.list_fields.return_value = ["FEE_SCHEDULE", "GROUPER_TYPE"]
            mock_fieldset.return_value = mock_fields

            result = one_to_n_funcs.methodology_breakout_secondary(
                methodology_dict, "Inpatient", "$500 per day",
                mock_fields, mock_fields, self.constants, self.filename
            )

        self.assertEqual(result["FEE_SCHEDULE"], "N/A")
        self.assertEqual(result["GROUPER_TYPE"], "N/A")

    # ==================== carveout_and_special_case Tests ====================

    @patch('src.investment.one_to_n_funcs.prompt_calls.prompt_carveout_check')
    @patch('src.investment.one_to_n_funcs.prompt_calls.prompt_special_case_breakout')
    def test_carveout_and_special_case_base_covered(self, mock_special_breakout, mock_carveout_check):
        """Test carveout check for BASE_COVERED_SERVICES."""
        mock_carveout_check.return_value = "BASE_COVERED_SERVICES"

        input_answers = [{"SERVICE_TERM": "Labs", "REIMB_TERM": "100% of Medicare"}]
        
        reimb_answers, special_answers = one_to_n_funcs.carveout_and_special_case(
            input_answers, self.constants, self.filename
        )

        self.assertEqual(len(reimb_answers), 1)
        self.assertEqual(len(special_answers), 0)
        self.assertEqual(reimb_answers[0]["CARVEOUT_CD"], "BASE_COVERED_SERVICES")
        self.assertEqual(reimb_answers[0]["CARVEOUT_IND"], "N")
        self.assertEqual(reimb_answers[0]["DEFAULT_IND"], "N")

    @patch('src.investment.one_to_n_funcs.prompt_calls.prompt_carveout_check')
    def test_carveout_and_special_case_default_term(self, mock_carveout_check):
        """Test carveout check for DEFAULT_TERM."""
        mock_carveout_check.return_value = "DEFAULT_TERM"

        input_answers = [{"SERVICE_TERM": "All other services", "REIMB_TERM": "As negotiated"}]
        
        reimb_answers, special_answers = one_to_n_funcs.carveout_and_special_case(
            input_answers, self.constants, self.filename
        )

        self.assertEqual(reimb_answers[0]["CARVEOUT_CD"], "DEFAULT_TERM")
        self.assertEqual(reimb_answers[0]["CARVEOUT_IND"], "N")
        self.assertEqual(reimb_answers[0]["DEFAULT_IND"], "Y")

    @patch('src.investment.one_to_n_funcs.prompt_calls.prompt_carveout_check')
    @patch('src.investment.one_to_n_funcs.prompt_calls.prompt_special_case_breakout')
    def test_carveout_and_special_case_trigger_cap(self, mock_special_breakout, mock_carveout_check):
        """Test carveout check for TRIGGER_CAP with breakout."""
        mock_carveout_check.return_value = "TRIGGER_CAP"
        mock_special_breakout.return_value = {"TRIGGER_AMOUNT": "$10000", "CAP_AMOUNT": "$50000"}

        input_answers = [{"SERVICE_TERM": "Surgery", "REIMB_TERM": "Trigger at $10k, cap at $50k"}]
        
        reimb_answers, special_answers = one_to_n_funcs.carveout_and_special_case(
            input_answers, self.constants, self.filename
        )

        self.assertEqual(reimb_answers[0]["CARVEOUT_CD"], "TRIGGER_CAP")
        self.assertEqual(reimb_answers[0]["CARVEOUT_IND"], "Y")
        self.assertEqual(reimb_answers[0]["DEFAULT_IND"], "N")
        self.assertIn("TRIGGER_AMOUNT", reimb_answers[0])

    @patch('src.investment.one_to_n_funcs.prompt_calls.prompt_carveout_check')
    def test_carveout_and_special_case_special_case_field(self, mock_carveout_check):
        """Test when result is a special case field, not a carveout."""
        mock_carveout_check.return_value = "OUTLIER_TERM"

        input_answers = [{"SERVICE_TERM": "High-cost cases", "REIMB_TERM": "Outlier payment above $25,000"}]
        
        reimb_answers, special_answers = one_to_n_funcs.carveout_and_special_case(
            input_answers, self.constants, self.filename
        )

        self.assertEqual(len(reimb_answers), 0)
        self.assertEqual(len(special_answers), 1)
        self.assertIn("OUTLIER_TERM", special_answers[0])
        self.assertNotIn("SERVICE_TERM", special_answers[0])
        self.assertNotIn("REIMB_TERM", special_answers[0])

    # ==================== special_case_breakout Tests ====================

    @patch('src.investment.one_to_n_funcs.prompt_calls.prompt_special_case_breakout')
    def test_special_case_breakout_outlier(self, mock_breakout):
        """Test special case breakout for outlier terms."""
        mock_breakout.return_value = {
            "OUTLIER_THRESHOLD": "$25000",
            "OUTLIER_PAYMENT_METHOD": "Cost-based"
        }

        input_answers = [{"OUTLIER_TERM": "Outlier payment for cases exceeding $25,000"}]
        
        with patch('src.investment.one_to_n_funcs.FieldSet') as mock_fieldset:
            mock_field = Mock()
            mock_field.get_breakout_template.return_value = "OUTLIER_BREAKOUT template"
            
            mock_fields = Mock()
            mock_fields.get_field.return_value = mock_field
            mock_fieldset.return_value = mock_fields

            result = one_to_n_funcs.special_case_breakout(input_answers, self.filename)

        self.assertEqual(len(result), 1)
        self.assertIn("OUTLIER_THRESHOLD", result[0])
        self.assertIn("OUTLIER_PAYMENT_METHOD", result[0])

    # ==================== filter_services_without_reimbursements Tests ====================

    @patch('src.investment.one_to_n_funcs.prompt_calls.validate_reimbursements_for_llm')
    def test_filter_services_all_valid(self, mock_validate):
        """Test filtering when all services have valid reimbursements."""
        mock_validate.return_value = True

        input_answers = [
            {"SERVICE_TERM": "Labs", "REIMB_TERM": "100% of Medicare"},
            {"SERVICE_TERM": "X-Ray", "REIMB_TERM": "Fee schedule"}
        ]

        result = one_to_n_funcs.filter_services_without_reimbursements(
            input_answers, self.filename
        )

        self.assertEqual(len(result), 2)
        self.assertEqual(mock_validate.call_count, 2)

    @patch('src.investment.one_to_n_funcs.prompt_calls.validate_reimbursements_for_llm')
    def test_filter_services_some_invalid(self, mock_validate):
        """Test filtering when some services lack reimbursements."""
        mock_validate.side_effect = [True, False, True]

        input_answers = [
            {"SERVICE_TERM": "Labs", "REIMB_TERM": "100% of Medicare"},
            {"SERVICE_TERM": "Definitions", "REIMB_TERM": "Standard eyeglass lenses as defined by DMAP"},
            {"SERVICE_TERM": "Surgery", "REIMB_TERM": "MS-DRG"}
        ]

        result = one_to_n_funcs.filter_services_without_reimbursements(
            input_answers, self.filename
        )

        self.assertEqual(len(result), 2)
        self.assertEqual(result[0]["SERVICE_TERM"], "Labs")
        self.assertEqual(result[1]["SERVICE_TERM"], "Surgery")

    @patch('src.investment.one_to_n_funcs.prompt_calls.validate_reimbursements_for_llm')
    def test_filter_services_all_invalid(self, mock_validate):
        """Test filtering when all services are invalid."""
        mock_validate.return_value = False

        input_answers = [
            {"SERVICE_TERM": "Admin", "REIMB_TERM": "Administrative terms apply"},
            {"SERVICE_TERM": "Definitions", "REIMB_TERM": "See glossary"}
        ]

        result = one_to_n_funcs.filter_services_without_reimbursements(
            input_answers, self.filename
        )

        self.assertEqual(len(result), 0)

    # ==================== split_reimb_dates Tests ====================

    @patch('src.investment.one_to_n_funcs.llm_utils.invoke_claude')
    @patch('src.investment.one_to_n_funcs.string_utils.universal_json_load')
    def test_split_reimb_dates_valid_range(self, mock_json_load, mock_invoke):
        """Test splitting valid date range."""
        mock_invoke.return_value = '{"start_date": "2024-01-01", "end_date": "2024-12-31"}'
        mock_json_load.return_value = {"start_date": "2024-01-01", "end_date": "2024-12-31"}

        input_data = [
            {"SERVICE_TERM": "Labs", "REIMB_DATES": "January 1, 2024 through December 31, 2024"}
        ]

        result = one_to_n_funcs.split_reimb_dates(input_data, self.filename)

        self.assertEqual(result[0]["REIMB_EFFECTIVE_DT"], "2024-01-01")
        self.assertEqual(result[0]["REIMB_TERMINATION_DT"], "2024-12-31")
        self.assertNotIn("REIMB_DATES", result[0])

    def test_split_reimb_dates_no_dates(self):
        """Test splitting when no REIMB_DATES field exists."""
        input_data = [
            {"SERVICE_TERM": "Labs", "REIMB_TERM": "100% of Medicare"}
        ]

        result = one_to_n_funcs.split_reimb_dates(input_data, self.filename)

        self.assertIsNone(result[0]["REIMB_EFFECTIVE_DT"])
        self.assertIsNone(result[0]["REIMB_TERMINATION_DT"])

    def test_split_reimb_dates_empty_dates(self):
        """Test splitting when REIMB_DATES is empty string."""
        input_data = [
            {"SERVICE_TERM": "Labs", "REIMB_DATES": ""}
        ]

        result = one_to_n_funcs.split_reimb_dates(input_data, self.filename)

        self.assertIsNone(result[0]["REIMB_EFFECTIVE_DT"])
        self.assertIsNone(result[0]["REIMB_TERMINATION_DT"])
        self.assertNotIn("REIMB_DATES", result[0])

    def test_split_reimb_dates_empty_input(self):
        """Test splitting with empty input list."""
        result = one_to_n_funcs.split_reimb_dates([], self.filename)
        self.assertEqual(result, [])

    def test_split_reimb_dates_none_input(self):
        """Test splitting with None input."""
        result = one_to_n_funcs.split_reimb_dates(None, self.filename)
        self.assertIsNone(result)

    @patch('src.investment.one_to_n_funcs.llm_utils.invoke_claude')
    def test_split_reimb_dates_llm_error(self, mock_invoke):
        """Test splitting when LLM call fails."""
        mock_invoke.side_effect = Exception("LLM API error")

        input_data = [
            {"SERVICE_TERM": "Labs", "REIMB_DATES": "January 1, 2024 through December 31, 2024"}
        ]

        result = one_to_n_funcs.split_reimb_dates(input_data, self.filename)

        # Should still remove REIMB_DATES and initialize new fields
        self.assertIsNone(result[0]["REIMB_EFFECTIVE_DT"])
        self.assertIsNone(result[0]["REIMB_TERMINATION_DT"])
        self.assertNotIn("REIMB_DATES", result[0])

    # ==================== get_lob_relationship Tests ====================

    @patch('src.investment.one_to_n_funcs.prompt_calls.prompt_lob_relationship')
    @patch('src.investment.one_to_n_funcs.string_utils.is_empty')
    def test_get_lob_relationship_with_program(self, mock_is_empty, mock_prompt_lob):
        """Test LOB-Program relationship extraction."""
        mock_is_empty.side_effect = [False, False, True]  # LOB not empty, PROGRAM not empty, PRODUCT empty
        mock_prompt_lob.return_value = "PROGRAM_SPECIFIC"

        input_data = [
            {"AARETE_DERIVED_LOB": "Medicaid", "AARETE_DERIVED_PROGRAM": "CHIP"}
        ]

        result = one_to_n_funcs.get_lob_relationship(
            input_data, self.exhibit_text, self.filename
        )

        self.assertEqual(result[0]["LOB_PROGRAM_RELATIONSHIP"], "PROGRAM_SPECIFIC")
        mock_prompt_lob.assert_called_once()

    @patch('src.investment.one_to_n_funcs.prompt_calls.prompt_lob_relationship')
    @patch('src.investment.one_to_n_funcs.string_utils.is_empty')
    def test_get_lob_relationship_with_product(self, mock_is_empty, mock_prompt_lob):
        """Test LOB-Product relationship extraction."""
        mock_is_empty.side_effect = [False, True, False]  # LOB not empty, PROGRAM empty, PRODUCT not empty
        mock_prompt_lob.return_value = "PRODUCT_SPECIFIC"

        input_data = [
            {"AARETE_DERIVED_LOB": "Commercial", "AARETE_DERIVED_PRODUCT": "HMO"}
        ]

        result = one_to_n_funcs.get_lob_relationship(
            input_data, self.exhibit_text, self.filename
        )

        self.assertEqual(result[0]["LOB_PRODUCT_RELATIONSHIP"], "PRODUCT_SPECIFIC")
        mock_prompt_lob.assert_called_once()

    @patch('src.investment.one_to_n_funcs.string_utils.is_empty')
    def test_get_lob_relationship_no_lob(self, mock_is_empty):
        """Test relationship extraction when LOB is empty."""
        mock_is_empty.return_value = True

        input_data = [
            {"AARETE_DERIVED_PROGRAM": "CHIP"}
        ]

        result = one_to_n_funcs.get_lob_relationship(
            input_data, self.exhibit_text, self.filename
        )

        # Should skip processing and return unchanged
        self.assertEqual(result, input_data)

    # ==================== one_to_n_cleaning Tests ====================

    @patch('src.investment.one_to_n_funcs.aarete_derived.get_crosswalk_fields')
    @patch('src.investment.one_to_n_funcs.get_lob_relationship')
    @patch('src.investment.one_to_n_funcs.aarete_derived.fill_na_mapping')
    @patch('src.investment.one_to_n_funcs.postprocessing_funcs.update_lob_for_duals')
    @patch('src.investment.one_to_n_funcs.split_reimb_dates')
    def test_one_to_n_cleaning_full_pipeline(self, mock_split_dates, mock_update_duals, 
                                              mock_fill_na, mock_get_lob, mock_crosswalk):
        """Test full one-to-n cleaning pipeline."""
        input_data = [{"SERVICE_TERM": "Labs", "REIMB_TERM": "100% of Medicare"}]
        
        mock_crosswalk.return_value = input_data
        mock_get_lob.return_value = input_data
        mock_fill_na.return_value = input_data
        mock_update_duals.return_value = input_data
        mock_split_dates.return_value = input_data

        result = one_to_n_funcs.one_to_n_cleaning(
            input_data, self.exhibit_text, self.constants, self.filename
        )

        # Verify all cleaning steps were called
        mock_crosswalk.assert_called_once()
        mock_get_lob.assert_called_once()
        mock_fill_na.assert_called_once()
        mock_update_duals.assert_called_once()
        mock_split_dates.assert_called_once()
        
        self.assertEqual(result, input_data)

    # ==================== lesser_of_distribution Tests ====================

    @patch('src.investment.one_to_n_funcs.prompt_calls.prompt_lesser_of_check')
    @patch('src.investment.one_to_n_funcs.prompt_calls.prompt_lesser_of_distribution')
    def test_lesser_of_distribution_with_lesser_term_not_overarching(self, mock_lesser_dist, mock_lesser_check):
        """Test lesser_of_distribution when 'less' is in term but not overarching."""
        mock_lesser_check.return_value = False  # Not overarching
        
        input_answers = [
            {"SERVICE_TERM": "Labs", "REIMB_TERM": "lesser of Medicare or billed charges"}
        ]
        
        result = one_to_n_funcs.lesser_of_distribution(
            input_answers, "exhibit text", "1", self.constants, self.filename
        )
        
        self.assertEqual(len(result), 1)
        self.assertEqual(result[0]["REIMB_TERM"], "lesser of Medicare or billed charges")
        mock_lesser_check.assert_called_once()
        mock_lesser_dist.assert_not_called()

    @patch('src.investment.one_to_n_funcs.prompt_calls.prompt_lesser_of_check')
    @patch('src.investment.one_to_n_funcs.prompt_calls.prompt_lesser_of_distribution')
    def test_lesser_of_distribution_with_lesser_term_overarching(self, mock_lesser_dist, mock_lesser_check):
        """Test lesser_of_distribution when 'less' is in term and is overarching."""
        mock_lesser_check.return_value = True  # Is overarching
        
        input_answers = [
            {"SERVICE_TERM": "Labs", "REIMB_TERM": "lesser of Medicare or billed charges"}
        ]
        
        result = one_to_n_funcs.lesser_of_distribution(
            input_answers, "exhibit text", "1", self.constants, self.filename
        )
        
        # When overarching, the answer should not be included (filtered out)
        self.assertEqual(len(result), 0)
        mock_lesser_check.assert_called_once()
        mock_lesser_dist.assert_not_called()

    @patch('src.investment.one_to_n_funcs.prompt_calls.prompt_lesser_of_check')
    @patch('src.investment.one_to_n_funcs.prompt_calls.prompt_lesser_of_distribution')
    def test_lesser_of_distribution_without_lesser_term(self, mock_lesser_dist, mock_lesser_check):
        """Test lesser_of_distribution when 'less' is not in term."""
        mock_lesser_dist.return_value = "100% of Medicare fee schedule"
        
        input_answers = [
            {"SERVICE_TERM": "Labs", "REIMB_TERM": "100% of Medicare"}
        ]
        
        result = one_to_n_funcs.lesser_of_distribution(
            input_answers, "exhibit text", "1", self.constants, self.filename
        )
        
        self.assertEqual(len(result), 1)
        self.assertEqual(result[0]["REIMB_TERM"], "100% of Medicare fee schedule")
        mock_lesser_check.assert_not_called()
        mock_lesser_dist.assert_called_once_with(
            "Labs", "100% of Medicare", "exhibit text", "1", self.filename
        )

    @patch('src.investment.one_to_n_funcs.prompt_calls.prompt_lesser_of_check')
    @patch('src.investment.one_to_n_funcs.prompt_calls.prompt_lesser_of_distribution')
    def test_lesser_of_distribution_multiple_answers(self, mock_lesser_dist, mock_lesser_check):
        """Test lesser_of_distribution with multiple reimbursement answers."""
        mock_lesser_check.return_value = False
        mock_lesser_dist.side_effect = [
            "Updated term 1",
            "Updated term 2"
        ]
        
        input_answers = [
            {"SERVICE_TERM": "Labs", "REIMB_TERM": "lesser of A or B"},
            {"SERVICE_TERM": "X-Ray", "REIMB_TERM": "Fee schedule"},
            {"SERVICE_TERM": "Surgery", "REIMB_TERM": "Per diem"}
        ]
        
        result = one_to_n_funcs.lesser_of_distribution(
            input_answers, "exhibit text", "1", self.constants, self.filename
        )
        
        self.assertEqual(len(result), 3)
        # First answer has "less" but not overarching, so kept as-is
        self.assertEqual(result[0]["REIMB_TERM"], "lesser of A or B")
        # Second and third answers get updated via prompt_lesser_of_distribution
        self.assertEqual(result[1]["REIMB_TERM"], "Updated term 1")
        self.assertEqual(result[2]["REIMB_TERM"], "Updated term 2")
        self.assertEqual(mock_lesser_check.call_count, 1)
        self.assertEqual(mock_lesser_dist.call_count, 2)

    @patch('src.investment.one_to_n_funcs.prompt_calls.prompt_lesser_of_check')
    @patch('src.investment.one_to_n_funcs.prompt_calls.prompt_lesser_of_distribution')
    def test_lesser_of_distribution_case_insensitive_less(self, mock_lesser_dist, mock_lesser_check):
        """Test that 'less' detection is case-insensitive."""
        mock_lesser_check.return_value = False
        
        input_answers = [
            {"SERVICE_TERM": "Labs", "REIMB_TERM": "LESSER of Medicare or billed"},
            {"SERVICE_TERM": "X-Ray", "REIMB_TERM": "Less of A or B"},
            {"SERVICE_TERM": "Surgery", "REIMB_TERM": "LESS of X or Y"}
        ]
        
        result = one_to_n_funcs.lesser_of_distribution(
            input_answers, "exhibit text", "1", self.constants, self.filename
        )
        
        # All three should trigger lesser_of_check (case-insensitive)
        self.assertEqual(mock_lesser_check.call_count, 3)
        self.assertEqual(len(result), 3)

    @patch('src.investment.one_to_n_funcs.prompt_calls.prompt_lesser_of_check')
    @patch('src.investment.one_to_n_funcs.prompt_calls.prompt_lesser_of_distribution')
    def test_lesser_of_distribution_empty_input(self, mock_lesser_dist, mock_lesser_check):
        """Test lesser_of_distribution with empty input list."""
        result = one_to_n_funcs.lesser_of_distribution(
            [], "exhibit text", "1", self.constants, self.filename
        )
        
        self.assertEqual(result, [])
        mock_lesser_check.assert_not_called()
        mock_lesser_dist.assert_not_called()

    @patch('src.investment.one_to_n_funcs.prompt_calls.prompt_lesser_of_check')
    @patch('src.investment.one_to_n_funcs.prompt_calls.prompt_lesser_of_distribution')
    def test_lesser_of_distribution_return_type(self, mock_lesser_dist, mock_lesser_check):
        """Test that lesser_of_distribution returns correct type structure."""
        mock_lesser_dist.return_value = "Updated term"
        
        input_answers = [
            {"SERVICE_TERM": "Labs", "REIMB_TERM": "100% of Medicare"}
        ]
        
        result = one_to_n_funcs.lesser_of_distribution(
            input_answers, "exhibit text", "1", self.constants, self.filename
        )
        
        # Verify return type is list[dict[str, str]]
        self.assertIsInstance(result, list)
        self.assertGreater(len(result), 0)
        self.assertIsInstance(result[0], dict)
        self.assertIn("SERVICE_TERM", result[0])
        self.assertIn("REIMB_TERM", result[0])


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