import json import re import unittest from unittest.mock import MagicMock, patch import pandas as pd import src.codes.code_funcs as code_funcs import src.prompts.prompt_templates as prompt_templates from src.constants.constants import Constants from src.constants.delimiters import Delimiter class TestCodeFuncs(unittest.TestCase): def setUp(self): """Set up test fixtures before each test method.""" # Clear the module-level valid code set caches between tests. code_funcs._VALID_PROC_CACHE.clear() code_funcs._VALID_REV_CACHE.clear() self.constants = MagicMock(spec=Constants) # Configure mock constants self.constants.SYNONYM_MAP = { "DME": "DURABLE MEDICAL EQUIPMENT", "SNF": "SKILLED NURSING FACILITY", } self.constants.REMOVAL_LIST = ["INPATIENT", "OUTPATIENT", "SERVICES"] self.constants.STOP_WORD_LIST = ["THE", "AND", "OF"] self.constants.DO_NOT_RUN = ["BY REPORT", "STOP-LOSS"] # Setup mappings self.constants.CPT_LEVEL1_MAPPING = {"12345": "Test Procedure"} self.constants.HCPCS_LEVEL1_MAPPING = {"J0001": "Test Drug"} self.constants.REV_LEVEL1_MAPPING = {"123": "Test Revenue"} self.constants.CPT_LEVEL2_MAPPING = {"67890": "Advanced Procedure"} self.constants.HCPCS_LEVEL2_MAPPING = {"J0002": "Advanced Drug"} self.constants.REV_MAPPING = {"0456": "Advanced Revenue"} # Full mappings for validation and description fill (Fix 1 / Fix 2) self.constants.CPT_FULL_MAPPING = { "12345": "Test Procedure", "99213": "Office visit, established patient", } self.constants.HCPCS_FULL_MAPPING = { "J0001": "Test Drug", "G0525": "Care management", } # Bill type mappings matching the actual JSON structure self.constants.BILL_TYPE_MAPPING = { "011X": "Inpatient Hospital", "012X": "Inpatient Hospital", "018X": "Inpatient Hospital", } # Reverse mapping returns a list of codes for each description self.constants.BILL_TYPE_REVERSE_MAPPING = { "Inpatient Hospital": ["011X", "012X", "018X"] } # Mock the embedding model and function mock_embedding = MagicMock() mock_embedding.index.search.return_value = ([0.9], [[0]]) mock_embedding.choices = ["Test Description"] self.constants.get_embedding.return_value = mock_embedding self.constants.EMBEDDING_MODEL = MagicMock() mock_encoded = MagicMock() mock_encoded.astype.return_value = "mock_vector" self.constants.EMBEDDING_MODEL.encode.return_value = mock_encoded def tearDown(self): """Clear prompt_templates constants cache so a MagicMock is not reused by other tests.""" prompt_templates.set_constants_for_instruction_cache(None) def test_clean_service(self): """Test the clean_service function with various inputs.""" # Test empty input self.assertEqual(code_funcs.clean_service("", self.constants), "") self.assertEqual(code_funcs.clean_service(None, self.constants), "") # Test synonym mapping self.assertEqual( code_funcs.clean_service("DME", self.constants), "DURABLE MEDICAL EQUIPMENT" ) # Test removal list self.assertEqual( code_funcs.clean_service("INPATIENT SERVICES", self.constants), "" ) # Test stop words only self.assertEqual(code_funcs.clean_service("THE AND OF", self.constants), "") # Test unlisted service self.assertEqual( code_funcs.clean_service("UNLISTED PROCEDURE", self.constants), "UNLISTED" ) # Test normal case self.assertEqual( code_funcs.clean_service("CT SCAN SNF", self.constants), "CT SCAN SKILLED NURSING FACILITY", ) @patch("src.utils.llm_utils.invoke_claude") @patch("src.utils.string_utils.universal_json_load") def test_code_explicit(self, mock_json_load, mock_invoke_claude): """Test the code_explicit function for extracting explicit codes.""" # Setup mocks mock_invoke_claude.return_value = '{"PROCEDURE_CD": "12345"}' mock_json_load.return_value = {"PROCEDURE_CD": "12345"} # Test successful extraction result = code_funcs.code_explicit( "TEST SERVICE", "TEST_METHODOLOGY", "test.pdf" ) self.assertEqual(result, {"PROCEDURE_CD": "12345"}) mock_invoke_claude.assert_called_once() @patch("src.utils.llm_utils.invoke_claude") @patch("src.utils.string_utils.universal_json_load") def test_code_category(self, mock_json_load, mock_invoke_claude): """Test the code_category function for detecting code categories.""" # Setup mocks mock_invoke_claude.return_value = '["Drug A", "Drug B"]' mock_json_load.return_value = ["Drug A", "Drug B"] # Setup test data proc_category = ["Category: J"] hcpcs_mapping = {"J0001": "Drug A", "J0002": "Drug B", "K0001": "Other Drug"} # Test category matching result = code_funcs.code_category( "J CODES", proc_category, hcpcs_mapping, "test.pdf" ) # Verify the results - code_category returns lists (JSON conversion happens in extract_codes_from_service) self.assertEqual(result["PROCEDURE_CD"], ["J0001", "J0002"]) self.assertEqual(result["PROCEDURE_CD_DESC"], ["Drug A", "Drug B"]) self.assertIn("J0001", result["PROCEDURE_CD"]) self.assertIn("J0002", result["PROCEDURE_CD"]) self.assertIn("Drug A", result["PROCEDURE_CD_DESC"]) self.assertIn("Drug B", result["PROCEDURE_CD_DESC"]) @patch("src.utils.llm_utils.invoke_claude") def test_code_implicit_special(self, mock_invoke_claude): """Test the code_implicit_special function for special case handling.""" # Setup mocks - parser returns a list mock_invoke_claude.return_value = '["Drugs"]' # Test drug category - returns a list result = code_funcs.code_implicit_special("DRUG SERVICE", "test.pdf") self.assertEqual(result["PROCEDURE_CD"], ["J0000-J9999"]) self.assertEqual(result["PROCEDURE_CD_DESC"], ["Drugs"]) # Test no match - parser returns empty list mock_invoke_claude.return_value = "[]" result = code_funcs.code_implicit_special("OTHER SERVICE", "test.pdf") self.assertEqual(result, {}) def test_get_embedding_levels(self): """Test the get_embedding_levels function for determining embedding levels.""" # Test level 1 with both codes enabled implicit_run_dict = {"PROCEDURE_CD": True, "REVENUE_CD": True} embedding_levels, cpt_mapping, hcpcs_mapping, rev_mapping = ( code_funcs.get_embedding_levels(1, implicit_run_dict, self.constants) ) self.assertIn("cpt_level1", embedding_levels) self.assertIn("hcpcs_level1", embedding_levels) self.assertIn("rev_level1", embedding_levels) self.assertEqual(cpt_mapping, self.constants.CPT_LEVEL1_MAPPING) # Test level 2 with only procedure code enabled implicit_run_dict = {"PROCEDURE_CD": True, "REVENUE_CD": False} embedding_levels, cpt_mapping, hcpcs_mapping, rev_mapping = ( code_funcs.get_embedding_levels(2, implicit_run_dict, self.constants) ) self.assertIn("cpt_level2", embedding_levels) self.assertIn("hcpcs_level2", embedding_levels) self.assertNotIn("rev", embedding_levels) self.assertEqual(cpt_mapping, self.constants.CPT_LEVEL2_MAPPING) def test_get_match_list(self): """Test the get_match_list function for retrieving matches from embeddings.""" # Configure mock for proper array-like return values mock_embedding = MagicMock() # Search returns a tuple of (scores, indices) where each is a 2D array mock_embedding.index.search.return_value = ( # Scores array - 2D array with shape [1, top_k] [[0.9]], # Indices array - 2D array with shape [1, top_k] [[0]], ) mock_embedding.choices = ["Test Description"] # Replace the mock embedding in constants for this test self.constants.get_embedding.return_value = mock_embedding # Test with constants mock configured in setUp match_list, highest_similarity = code_funcs.get_match_list( "TEST SERVICE", ["cpt_level1"], self.constants, 1 ) # Verify results self.assertEqual(match_list, ["Test Description"]) self.assertEqual(highest_similarity, 0.9) # Verify the appropriate method calls self.constants.EMBEDDING_MODEL.encode.assert_called_once() self.constants.get_embedding.assert_called_once_with("cpt_level1") @patch("src.utils.llm_utils.invoke_claude") @patch("src.utils.string_utils.universal_json_load") @patch("src.codes.code_funcs.get_match_list") def test_code_implicit_rag( self, mock_get_match_list, mock_json_load, mock_invoke_claude ): """Test the code_implicit_rag function for RAG-based code extraction.""" # Mock get_match_list to return properly formatted results mock_get_match_list.return_value = (["Test Procedure"], 0.9) # Setup mocks for Claude responses mock_invoke_claude.return_value = '["Test Procedure"]' mock_json_load.return_value = ["Test Procedure"] # Test with successful match implicit_run_dict = {"PROCEDURE_CD": True, "REVENUE_CD": True} # Patch the get_embedding_levels function to return predictable values with patch( "src.codes.code_funcs.get_embedding_levels" ) as mock_get_embedding_levels: # Configure mock to return appropriate values for embedding levels mock_get_embedding_levels.return_value = ( ["cpt_level1"], {"12345": "Test Procedure"}, {}, {}, ) # Call the function result = code_funcs.code_implicit_rag( "TEST SERVICE", implicit_run_dict, "test.pdf", self.constants ) # Verify results - code_implicit_rag returns list of candidates with source Level 1/Level 2 self.assertIsInstance(result, list) self.assertGreaterEqual(len(result), 1) codes_found = False for item in result: self.assertIn("source", item) self.assertIn("code_answer_dict", item) d = item["code_answer_dict"] if d.get("PROCEDURE_CD") == ["12345"]: codes_found = True self.assertEqual(d["PROCEDURE_CD_DESC"], ["Test Procedure"]) break self.assertTrue( codes_found, "Expected PROCEDURE_CD ['12345'] in a candidate" ) # Test exception handling - returns list (possibly empty) mock_json_load.side_effect = Exception("Test error") result = code_funcs.code_implicit_rag( "TEST SERVICE", implicit_run_dict, "test.pdf", self.constants ) self.assertIsInstance(result, list) @patch("src.utils.llm_utils.invoke_claude") @patch("src.utils.string_utils.universal_json_load") @patch("src.codes.code_funcs.get_match_list") def test_code_implicit_rag_invalid_service_short_circuits( self, mock_get_match_list, mock_json_load, mock_invoke_claude ): """INVALID_SERVICE in RAG response is a hard stop: no codes added.""" mock_get_match_list.return_value = (["Test Procedure"], 0.9) mock_invoke_claude.return_value = '["INVALID_SERVICE", "Test Procedure"]' mock_json_load.return_value = ["INVALID_SERVICE", "Test Procedure"] implicit_run_dict = {"PROCEDURE_CD": True, "REVENUE_CD": True} with patch( "src.codes.code_funcs.get_embedding_levels" ) as mock_get_embedding_levels: mock_get_embedding_levels.return_value = ( ["cpt_level1"], {"12345": "Test Procedure"}, {}, {}, ) result = code_funcs.code_implicit_rag( "TEST SERVICE", implicit_run_dict, "test.pdf", self.constants ) # INVALID_SERVICE yields no candidates with codes self.assertIsInstance(result, list) self.assertEqual(len(result), 0) @patch("src.utils.llm_utils.invoke_claude") def test_fill_bill_type_list_reverse_mapping(self, mock_invoke_claude): """fill_bill_type handles list values in BILL_TYPE_REVERSE_MAPPING.""" mock_invoke_claude.return_value = '["Inpatient Hospital"]' reverse_mapping = {"Inpatient Hospital": ["11X", "12X"]} answer_dict = {} result = code_funcs.fill_bill_type( "INPATIENT SERVICE", answer_dict, self.constants.BILL_TYPE_MAPPING, reverse_mapping, ) self.assertEqual(result["BILL_TYPE_CD"], ["11X", "12X"]) @patch("src.utils.llm_utils.invoke_claude") def test_code_last_check(self, mock_invoke_claude): """Test the code_last_check function for categorizing non-matched services.""" # Parser is format-aware and returns 'Specific' or 'Generic' directly. # Test specific case mock_invoke_claude.return_value = '["Specific"]' result = code_funcs.code_last_check("SPECIAL SERVICE", "test.pdf") self.assertEqual(result, "Specific") # Test generic case mock_invoke_claude.return_value = '["Generic"]' result = code_funcs.code_last_check("GENERIC SERVICE", "test.pdf") self.assertEqual(result, "Generic") def test_format_implicit_candidate_for_prompt(self): """Test _format_implicit_candidate_for_prompt formats a candidate for arbitration.""" candidate = { "source": "Category", "code_answer_dict": { "PROCEDURE_CD": ["J0001", "J0002"], "PROCEDURE_CD_DESC": ["Drug A", "Drug B"], }, } result = code_funcs._format_implicit_candidate_for_prompt(candidate, 0) self.assertIn("Candidate 0 (Source: Category)", result) self.assertIn("PROCEDURE_CD", result) self.assertIn("J0001", result) @patch("src.codes.code_funcs.code_implicit_rag") @patch("src.codes.code_funcs.code_implicit_special") @patch("src.codes.code_funcs.code_category") def test_build_implicit_candidates_empty_when_none_return( self, mock_category, mock_special, mock_rag ): """Test build_implicit_candidates returns empty list when no strategy returns.""" mock_category.return_value = {} mock_special.return_value = {} mock_rag.return_value = [] # RAG returns list of candidates implicit_run_dict = {"PROCEDURE_CD": True, "REVENUE_CD": True} result = code_funcs.build_implicit_candidates( "SERVICE", {}, implicit_run_dict, "test.pdf", self.constants, ) self.assertEqual(result, []) @patch("src.codes.code_funcs.code_implicit_rag") @patch("src.codes.code_funcs.code_implicit_special") @patch("src.codes.code_funcs.code_category") def test_build_implicit_candidates_category_when_explicit_has_category( self, mock_category, mock_special, mock_rag ): """Test build_implicit_candidates runs category and adds candidate when explicit had Category.""" mock_category.return_value = { "PROCEDURE_CD": ["J0001"], "PROCEDURE_CD_DESC": ["Drug A"], } mock_special.return_value = {} mock_rag.return_value = [] # RAG returns list of candidates implicit_run_dict = {"PROCEDURE_CD": True, "REVENUE_CD": True} code_after_explicit = {"PROCEDURE_CD": ["Category: J"]} result = code_funcs.build_implicit_candidates( "J CODES", code_after_explicit, implicit_run_dict, "test.pdf", self.constants, ) self.assertEqual(len(result), 1) self.assertEqual(result[0]["source"], "Category") self.assertEqual(result[0]["code_answer_dict"]["PROCEDURE_CD"], ["J0001"]) mock_category.assert_called_once() @patch("src.codes.code_funcs.code_implicit_rag") @patch("src.codes.code_funcs.code_implicit_special") @patch("src.codes.code_funcs.code_category") def test_build_implicit_candidates_skips_category_without_category_in_explicit( self, mock_category, mock_special, mock_rag ): """Test build_implicit_candidates does not run category when explicit has no Category.""" mock_special.return_value = { "PROCEDURE_CD": ["J0000-J9999"], "PROCEDURE_CD_DESC": ["Drugs"], } mock_rag.return_value = [] # RAG returns list of candidates implicit_run_dict = {"PROCEDURE_CD": True, "REVENUE_CD": True} result = code_funcs.build_implicit_candidates( "DRUG SERVICE", {}, implicit_run_dict, "test.pdf", self.constants, ) mock_category.assert_not_called() self.assertEqual(len(result), 1) self.assertEqual(result[0]["source"], "Special") @patch("src.codes.code_funcs.code_implicit_rag") @patch("src.codes.code_funcs.code_implicit_special") @patch("src.codes.code_funcs.code_category") def test_build_implicit_candidates_collects_category_special_rag( self, mock_category, mock_special, mock_rag ): """Test build_implicit_candidates collects candidates from all three when all return.""" mock_category.return_value = { "PROCEDURE_CD": ["A0000-A9999"], "PROCEDURE_CD_DESC": ["A range"], } mock_special.return_value = { "PROCEDURE_CD": ["J0000-J9999"], "PROCEDURE_CD_DESC": ["Drugs"], } mock_rag.return_value = [ { "source": "Level 2", "code_answer_dict": { "PROCEDURE_CD": ["99213"], "PROCEDURE_CD_DESC": ["Office visit"], "CODE_METHODOLOGY": "Implicit - Level 2", }, }, ] implicit_run_dict = {"PROCEDURE_CD": True, "REVENUE_CD": True} code_after_explicit = {"PROCEDURE_CD": ["Category: A"]} result = code_funcs.build_implicit_candidates( "SERVICE", code_after_explicit, implicit_run_dict, "test.pdf", self.constants, ) self.assertEqual(len(result), 3) sources = [c["source"] for c in result] self.assertIn("Category", sources) self.assertIn("Special", sources) self.assertIn("Level 2", sources) def test_code_implicit_arbitration_zero_candidates_returns_none(self): """Test code_implicit_arbitration returns None when candidates list is empty.""" result = code_funcs.code_implicit_arbitration("SERVICE", [], "test.pdf") self.assertIsNone(result) def test_code_implicit_arbitration_single_candidate_returns_it_no_llm(self): """Test code_implicit_arbitration with one candidate returns it without calling LLM.""" candidates = [ { "source": "Category", "code_answer_dict": { "PROCEDURE_CD": ["J0001"], "PROCEDURE_CD_DESC": ["Drug A"], }, } ] result = code_funcs.code_implicit_arbitration("J CODES", candidates, "test.pdf") self.assertIsNotNone(result) self.assertEqual(result["PROCEDURE_CD"], ["J0001"]) self.assertEqual( result["CODE_METHODOLOGY"], "Implicit - Arbitration (Category)", ) @patch("src.utils.llm_utils.invoke_claude") def test_code_implicit_arbitration_multiple_candidates_uses_chosen_index( self, mock_invoke_claude ): """Test code_implicit_arbitration with 2+ candidates calls LLM and returns chosen candidate.""" mock_invoke_claude.return_value = '{"chosen_index": 1}' candidates = [ { "source": "Category", "code_answer_dict": {"PROCEDURE_CD": ["A0000-A9999"]}, }, {"source": "Level 2", "code_answer_dict": {"PROCEDURE_CD": ["99213"]}}, ] result = code_funcs.code_implicit_arbitration( "Office visit", candidates, "test.pdf" ) self.assertIsNotNone(result) self.assertEqual(result["PROCEDURE_CD"], ["99213"]) self.assertEqual( result["CODE_METHODOLOGY"], "Implicit - Arbitration (Level 2)", ) mock_invoke_claude.assert_called_once() @patch("src.utils.llm_utils.invoke_claude") def test_code_implicit_arbitration_no_match_returns_none(self, mock_invoke_claude): """Test code_implicit_arbitration returns None when LLM returns no_match.""" mock_invoke_claude.return_value = '{"no_match": true}' candidates = [ { "source": "Category", "code_answer_dict": {"PROCEDURE_CD": ["A0000-A9999"]}, }, { "source": "Special", "code_answer_dict": {"PROCEDURE_CD": ["10004-69990"]}, }, ] result = code_funcs.code_implicit_arbitration( "Unclear service", candidates, "test.pdf" ) self.assertIsNone(result) @patch("src.utils.llm_utils.invoke_claude") def test_code_implicit_arbitration_invalid_index_fallback_to_zero( self, mock_invoke_claude ): """Test code_implicit_arbitration falls back to index 0 when chosen_index out of range.""" mock_invoke_claude.return_value = '{"chosen_index": 99}' candidates = [ {"source": "Category", "code_answer_dict": {"PROCEDURE_CD": ["J0001"]}}, ] result = code_funcs.code_implicit_arbitration("SERVICE", candidates, "test.pdf") self.assertIsNotNone(result) self.assertEqual(result["PROCEDURE_CD"], ["J0001"]) self.assertEqual( result["CODE_METHODOLOGY"], "Implicit - Arbitration (Category)", ) @patch("src.utils.llm_utils.invoke_claude") def test_fill_bill_type(self, mock_invoke_claude): """Test the fill_bill_type function for populating bill type information.""" mock_invoke_claude.return_value = '["Inpatient Hospital"]' # Test with valid bill type answer_dict = {} result = code_funcs.fill_bill_type( "INPATIENT SERVICE", answer_dict, self.constants.BILL_TYPE_MAPPING, self.constants.BILL_TYPE_REVERSE_MAPPING, ) # Should return all codes for "Inpatient Hospital" description # Reverse mapping returns ["011X", "012X", "018X"] for "Inpatient Hospital" self.assertEqual(result["BILL_TYPE_CD"], ["011X", "012X", "018X"]) self.assertEqual(result["BILL_TYPE_CD_DESC"], ["Inpatient Hospital"]) # Test with empty response mock_invoke_claude.return_value = "[]" answer_dict = {"SERVICE_TERM": "Test"} result = code_funcs.fill_bill_type( "OTHER SERVICE", answer_dict, self.constants.BILL_TYPE_MAPPING, self.constants.BILL_TYPE_REVERSE_MAPPING, ) self.assertEqual(result, {"SERVICE_TERM": "Test"}) def test_get_implicit_runs(self): """Test the get_implicit_runs function for determining which code runs to execute.""" # Test hospital claim type result = code_funcs.get_implicit_runs({"AARETE_DERIVED_CLAIM_TYPE_CD": "H"}) self.assertTrue(result["REVENUE_CD"]) # Test medical claim type with bill type result = code_funcs.get_implicit_runs( {"AARETE_DERIVED_CLAIM_TYPE_CD": "M", "BILL_TYPE_CD_DESC": "Outpatient"} ) self.assertTrue(result["REVENUE_CD"]) self.assertTrue(result["PROCEDURE_CD"]) # Test inpatient bill type result = code_funcs.get_implicit_runs( { "AARETE_DERIVED_CLAIM_TYPE_CD": "M", "BILL_TYPE_CD_DESC": "Inpatient Hospital", } ) self.assertTrue(result["REVENUE_CD"]) self.assertFalse(result["PROCEDURE_CD"]) @patch("src.codes.code_funcs.clean_service") @patch("src.codes.code_funcs.fill_bill_type") @patch("src.codes.code_funcs.get_implicit_runs") @patch("src.codes.code_funcs.code_explicit") @patch("src.codes.code_funcs.code_category") @patch("src.codes.code_funcs.code_implicit_special") @patch("src.codes.code_funcs.code_implicit_rag") @patch("src.codes.code_funcs.code_last_check") def test_extract_codes_from_service( self, mock_last_check, mock_implicit_rag, mock_implicit_special, mock_category, mock_explicit, mock_implicit_runs, mock_fill_bill_type, mock_clean_service, ): """Test the extract_codes_from_service function end-to-end.""" # Setup mocks mock_clean_service.return_value = "CLEAN SERVICE" mock_fill_bill_type.return_value = { "SERVICE_TERM": "TEST SERVICE", "BILL_TYPE_CD_DESC": "Outpatient", } mock_implicit_runs.return_value = {"PROCEDURE_CD": True, "REVENUE_CD": True} # Test explicit code match - extract_codes_from_service normalizes to JSON list format mock_explicit.return_value = {"PROCEDURE_CD": "12345"} result = code_funcs.extract_codes_from_service( {"SERVICE_TERM": "TEST SERVICE"}, self.constants ) self.assertEqual(result["PROCEDURE_CD"], '["12345"]') self.assertEqual(result["CODE_METHODOLOGY"], "Explicit") # Test unlisted service mock_clean_service.return_value = "UNLISTED" result = code_funcs.extract_codes_from_service( {"SERVICE_TERM": "UNLISTED SERVICE"}, self.constants ) self.assertEqual(result["PROCEDURE_CD_DESC"], "UNLISTED") # Test category match (arbitration with single candidate returns it) mock_clean_service.return_value = "CLEAN SERVICE" mock_explicit.return_value = {"PROCEDURE_CD": ["Category: J"]} mock_category.return_value = { "PROCEDURE_CD": ["J0001"], "PROCEDURE_CD_DESC": ["Drug A"], } result = code_funcs.extract_codes_from_service( {"SERVICE_TERM": "J CODES"}, self.constants ) self.assertEqual( result["CODE_METHODOLOGY"], "Implicit - Arbitration (Category)" ) # Test special case match (arbitration with single candidate) mock_explicit.return_value = {} mock_implicit_special.return_value = { "PROCEDURE_CD": "J0000-J9999", "PROCEDURE_CD_DESC": "Drugs", } result = code_funcs.extract_codes_from_service( {"SERVICE_TERM": "DRUG SERVICE"}, self.constants ) # Range J0000-J9999 not in mock mapping so methodology gets - Unmatched self.assertIn("Implicit - Arbitration (Special)", result["CODE_METHODOLOGY"]) self.assertIn("Unmatched", result["CODE_METHODOLOGY"]) # Test RAG match (arbitration with single candidate; RAG returns list with Level 1) mock_implicit_special.return_value = {} mock_implicit_rag.return_value = [ { "source": "Level 1", "code_answer_dict": { "PROCEDURE_CD": "12345", "CODE_METHODOLOGY": "Implicit - Level 1", }, }, ] result = code_funcs.extract_codes_from_service( {"SERVICE_TERM": "TEST SERVICE"}, self.constants ) self.assertEqual(result["CODE_METHODOLOGY"], "Implicit - Arbitration (Level 1)") # Test no match - specific mock_implicit_rag.return_value = {} mock_last_check.return_value = "Specific" result = code_funcs.extract_codes_from_service( {"SERVICE_TERM": "TEST SERVICE"}, self.constants ) self.assertEqual(result["CODE_METHODOLOGY"], "Specific - No Match") # Test no match - generic mock_last_check.return_value = "Generic" result = code_funcs.extract_codes_from_service( {"SERVICE_TERM": "TEST SERVICE"}, self.constants ) self.assertEqual(result["CODE_METHODOLOGY"], "Generic - No Match") def test_code_in_range(self): """Test _code_in_range for numeric and alphanumeric range keys.""" self.assertTrue(code_funcs._code_in_range("00100-01999", "01234")) self.assertTrue(code_funcs._code_in_range("00100-01999", "00100")) self.assertTrue(code_funcs._code_in_range("00100-01999", "01999")) self.assertFalse(code_funcs._code_in_range("00100-01999", "99999")) self.assertFalse(code_funcs._code_in_range("00100-01999", "02000")) self.assertTrue(code_funcs._code_in_range("A0021-A0999", "A0500")) self.assertFalse(code_funcs._code_in_range("A0021-A0999", "B0001")) self.assertFalse(code_funcs._code_in_range("", "12345")) self.assertFalse(code_funcs._code_in_range("00100-01999", "")) def test_lookup_procedure_description(self): """Test _lookup_procedure_description: exact match and range match.""" desc = code_funcs._lookup_procedure_description("12345", self.constants) self.assertEqual(desc, "Test Procedure") desc = code_funcs._lookup_procedure_description("99213", self.constants) self.assertEqual(desc, "Office visit, established patient") desc = code_funcs._lookup_procedure_description( "NOT_ESTABLISHED", self.constants ) self.assertEqual(desc, "") desc = code_funcs._lookup_procedure_description("UNMAPPED99", self.constants) self.assertEqual(desc, "") def test_lookup_revenue_description(self): """Test _lookup_revenue_description.""" desc = code_funcs._lookup_revenue_description("0456", self.constants) self.assertEqual(desc, "Advanced Revenue") desc = code_funcs._lookup_revenue_description("9999", self.constants) self.assertEqual(desc, "") def test_fill_code_descriptions_from_mappings(self): """Test fill_code_descriptions_from_mappings fills _DESC from mappings.""" answer_dict = { "PROCEDURE_CD": ["12345", "99213"], "REVENUE_CD": ["0456"], } code_funcs.fill_code_descriptions_from_mappings(answer_dict, self.constants) self.assertEqual( answer_dict["PROCEDURE_CD_DESC"], ["Test Procedure", "Office visit, established patient"], ) self.assertEqual(answer_dict["REVENUE_CD_DESC"], ["Advanced Revenue"]) def test_fill_code_descriptions_leave_unmapped_unchanged(self): """Test fill leaves _DESC empty for unmapped codes and NOT_ESTABLISHED.""" answer_dict = { "PROCEDURE_CD": ["12345", "NOT_ESTABLISHED", "UNMAPPED99"], "PROCEDURE_CD_DESC": ["", "", ""], } code_funcs.fill_code_descriptions_from_mappings(answer_dict, self.constants) self.assertEqual(answer_dict["PROCEDURE_CD_DESC"][0], "Test Procedure") self.assertEqual(answer_dict["PROCEDURE_CD_DESC"][1], "") self.assertEqual(answer_dict["PROCEDURE_CD_DESC"][2], "") def test_validate_explicit_codes_keeps_valid_filters_invalid(self): """Format-invalid codes dropped; format-valid preserved; has_unmatched if not in mapping.""" code_answer_dict = { "PROCEDURE_CD": ["12345", "99213", "garbage7"], "REVENUE_CD": ["0456", "9999"], } filtered, has_unmatched = code_funcs.validate_explicit_codes( code_answer_dict, self.constants ) self.assertTrue(has_unmatched) self.assertEqual(filtered["PROCEDURE_CD"], ["12345", "99213"]) self.assertEqual(filtered["REVENUE_CD"], ["0456", "9999"]) def test_validate_explicit_codes_keeps_category_and_not_established(self): """Test validate_explicit_codes keeps Category: and NOT_ESTABLISHED; no unmatched.""" code_answer_dict = {"PROCEDURE_CD": ["Category: J", "NOT_ESTABLISHED"]} filtered, has_unmatched = code_funcs.validate_explicit_codes( code_answer_dict, self.constants ) self.assertFalse(has_unmatched) self.assertEqual(filtered["PROCEDURE_CD"], ["Category: J", "NOT_ESTABLISHED"]) def test_validate_explicit_codes_all_valid(self): """Test validate_explicit_codes returns has_unmatched False when all in mapping.""" code_answer_dict = {"PROCEDURE_CD": ["12345"], "REVENUE_CD": ["0456"]} filtered, has_unmatched = code_funcs.validate_explicit_codes( code_answer_dict, self.constants ) self.assertFalse(has_unmatched) self.assertEqual(filtered["PROCEDURE_CD"], ["12345"]) self.assertEqual(filtered["REVENUE_CD"], ["0456"]) def test_validate_explicit_codes_multiple_and_range(self): """Range key that is an explicit key in mapping is kept and counts as mapped.""" self.constants.CPT_LEVEL1_MAPPING["00100-01999"] = "Anesthesia" code_answer_dict = { "PROCEDURE_CD": ["12345", "99213", "00100-01999"], "REVENUE_CD": ["0456"], } filtered, has_unmatched = code_funcs.validate_explicit_codes( code_answer_dict, self.constants ) self.assertFalse(has_unmatched) self.assertEqual(filtered["PROCEDURE_CD"], ["12345", "99213", "00100-01999"]) self.assertEqual(filtered["REVENUE_CD"], ["0456"]) def test_validate_explicit_codes_range_not_in_mapping_kept_tagged_unmatched(self): """Format-valid range not in mapping is kept; has_unmatched True for tag.""" # Mock has 99213 in CPT_FULL_MAPPING but not 10021-69990; range is kept and tagged code_answer_dict = { "PROCEDURE_CD": ["99213", "10021-69990"], "REVENUE_CD": ["0456"], } filtered, has_unmatched = code_funcs.validate_explicit_codes( code_answer_dict, self.constants ) self.assertTrue(has_unmatched) self.assertEqual(filtered["PROCEDURE_CD"], ["99213", "10021-69990"]) self.assertEqual(filtered["REVENUE_CD"], ["0456"]) def test_validate_explicit_codes_single_code_unmapped(self): """Format-valid code not in exact mapping is preserved and has_unmatched True.""" code_answer_dict = {"PROCEDURE_CD": ["01234"], "REVENUE_CD": []} filtered, has_unmatched = code_funcs.validate_explicit_codes( code_answer_dict, self.constants ) self.assertTrue(has_unmatched) self.assertEqual(filtered["PROCEDURE_CD"], ["01234"]) def test_validate_explicit_codes_rejects_icd10_pcs_in_procedure(self): """ICD-10-PCS (7-char) filtered from PROCEDURE_CD; CPT/HCPCS format kept; S2900 in mapping.""" self.constants.HCPCS_FULL_MAPPING["S2900"] = "Robotic surgical system" code_answer_dict = { "PROCEDURE_CD": ["0E0V0CZ", "0E0W0CZ", "S2900"], "REVENUE_CD": [], } filtered, has_unmatched = code_funcs.validate_explicit_codes( code_answer_dict, self.constants ) self.assertFalse(has_unmatched) self.assertEqual(filtered["PROCEDURE_CD"], ["S2900"]) def test_has_unmatched_codes_true_for_unmapped_procedure(self): """_has_unmatched_codes returns True when procedure code is format-valid but not in mapping.""" d = {"PROCEDURE_CD": ["99999"], "REVENUE_CD": []} self.assertTrue(code_funcs._has_unmatched_codes(d, self.constants)) def test_has_unmatched_codes_true_for_unmapped_revenue(self): """_has_unmatched_codes returns True when revenue code is format-valid but not in mapping.""" d = {"PROCEDURE_CD": [], "REVENUE_CD": ["9999"]} self.assertTrue(code_funcs._has_unmatched_codes(d, self.constants)) def test_has_unmatched_codes_false_when_all_mapped(self): """_has_unmatched_codes returns False when all procedure/revenue codes are in mapping.""" d = {"PROCEDURE_CD": ["99213"], "REVENUE_CD": ["0456"]} self.assertFalse(code_funcs._has_unmatched_codes(d, self.constants)) @patch("src.codes.code_funcs.code_explicit") def test_extract_codes_explicit_unmatched(self, mock_explicit): """Explicit path sets Unmatched when any format-valid code not in mapping; no retry.""" mock_explicit.return_value = { "PROCEDURE_CD": ["12345", "99999"], "REVENUE_CD": [], } with patch("src.codes.code_funcs.clean_service", return_value="CLEAN"): with patch("src.codes.code_funcs.fill_bill_type", return_value={}): with patch( "src.codes.code_funcs.get_implicit_runs", return_value={"PROCEDURE_CD": True, "REVENUE_CD": True}, ): result = code_funcs.extract_codes_from_service( {"SERVICE_TERM": "TEST"}, self.constants ) self.assertEqual(result["CODE_METHODOLOGY"], "Explicit - Unmatched") self.assertIn("12345", str(result["PROCEDURE_CD"])) self.assertIn("99999", str(result["PROCEDURE_CD"])) self.assertEqual(mock_explicit.call_count, 1) @patch("src.codes.code_funcs.code_explicit") def test_extract_codes_explicit_description_filled(self, mock_explicit): """Test explicit path fills PROCEDURE_CD_DESC from mapping.""" mock_explicit.return_value = { "PROCEDURE_CD": ["99213"], "REVENUE_CD": ["0456"], } with patch("src.codes.code_funcs.clean_service", return_value="CLEAN"): with patch("src.codes.code_funcs.fill_bill_type", return_value={}): with patch( "src.codes.code_funcs.get_implicit_runs", return_value={"PROCEDURE_CD": True, "REVENUE_CD": True}, ): result = code_funcs.extract_codes_from_service( {"SERVICE_TERM": "Office visit"}, self.constants ) self.assertEqual(result["CODE_METHODOLOGY"], "Explicit") self.assertIn("99213", str(result["PROCEDURE_CD"])) self.assertIn("Office visit", str(result["PROCEDURE_CD_DESC"])) self.assertIn("Advanced Revenue", str(result["REVENUE_CD_DESC"])) @patch("src.codes.code_funcs.code_implicit_arbitration") @patch("src.codes.code_funcs.build_implicit_candidates") @patch("src.codes.code_funcs.code_explicit") def test_extract_codes_implicit_arbitration_unmatched( self, mock_explicit, mock_build_candidates, mock_arbitration ): """Implicit path appends ' - Unmatched' when chosen candidate has unmapped code.""" mock_explicit.return_value = { "PROCEDURE_CD": ["Category: J"], "REVENUE_CD": [], } mock_build_candidates.return_value = [ { "source": "RAG", "code_answer_dict": { "PROCEDURE_CD": ["99999"], "REVENUE_CD": [], }, } ] chosen = { "PROCEDURE_CD": ["99999"], "REVENUE_CD": [], "CODE_METHODOLOGY": "Implicit - Arbitration (RAG)", } mock_arbitration.return_value = chosen with patch("src.codes.code_funcs.clean_service", return_value="CLEAN"): with patch("src.codes.code_funcs.fill_bill_type", return_value={}): with patch( "src.codes.code_funcs.get_implicit_runs", return_value={"PROCEDURE_CD": True, "REVENUE_CD": True}, ): result = code_funcs.extract_codes_from_service( {"SERVICE_TERM": "J CODES"}, self.constants ) self.assertEqual( result["CODE_METHODOLOGY"], "Implicit - Arbitration (RAG) - Unmatched" ) self.assertIn("99999", str(result["PROCEDURE_CD"])) @patch("src.codes.code_funcs.code_implicit_arbitration") @patch("src.codes.code_funcs.build_implicit_candidates") @patch("src.codes.code_funcs.code_explicit") def test_extract_codes_implicit_arbitration_all_mapped( self, mock_explicit, mock_build_candidates, mock_arbitration ): """Implicit path does not append ' - Unmatched' when all codes are in mapping.""" mock_explicit.return_value = { "PROCEDURE_CD": ["Category: J"], "REVENUE_CD": [], } mock_build_candidates.return_value = [ { "source": "RAG", "code_answer_dict": { "PROCEDURE_CD": ["99213"], "REVENUE_CD": ["0456"], }, } ] chosen = { "PROCEDURE_CD": ["99213"], "REVENUE_CD": ["0456"], "CODE_METHODOLOGY": "Implicit - Arbitration (RAG)", } mock_arbitration.return_value = chosen with patch("src.codes.code_funcs.clean_service", return_value="CLEAN"): with patch("src.codes.code_funcs.fill_bill_type", return_value={}): with patch( "src.codes.code_funcs.get_implicit_runs", return_value={"PROCEDURE_CD": True, "REVENUE_CD": True}, ): result = code_funcs.extract_codes_from_service( {"SERVICE_TERM": "J CODES"}, self.constants ) self.assertEqual(result["CODE_METHODOLOGY"], "Implicit - Arbitration (RAG)") self.assertIn("99213", str(result["PROCEDURE_CD"])) self.assertIn("0456", str(result["REVENUE_CD"])) @patch("src.codes.code_funcs.extract_codes_from_service") def test_code_breakout(self, mock_extract_codes): """Test the code_breakout function for processing DataFrame records.""" # BILL_TYPE_CD_DESC set so pre-pass skips fill_bill_type (avoids LLM on CI). test_df = pd.DataFrame( [ { "SERVICE_TERM": "SERVICE A", "BILL_TYPE_CD_DESC": "Inpatient Hospital", }, { "SERVICE_TERM": "SERVICE B", "BILL_TYPE_CD_DESC": "Inpatient Hospital", }, ] ) # Return value by SERVICE_TERM so order is correct under ThreadPoolExecutor. def return_by_service(rep, constants): st = rep.get("SERVICE_TERM", "") if st == "SERVICE A": return {"SERVICE_TERM": "SERVICE A", "PROCEDURE_CD": ["12345"]} return {"SERVICE_TERM": "SERVICE B", "PROCEDURE_CD": ["67890"]} mock_extract_codes.side_effect = return_by_service # Test function result_df = code_funcs.code_breakout(test_df, self.constants) # Verify results - code_breakout normalizes to JSON list format self.assertIsInstance(result_df, pd.DataFrame) self.assertEqual(len(result_df), 2) self.assertEqual(result_df.iloc[0]["PROCEDURE_CD"], '["12345"]') self.assertEqual(result_df.iloc[1]["PROCEDURE_CD"], '["67890"]') # Verify mock calls self.assertEqual(mock_extract_codes.call_count, 2) if __name__ == "__main__": unittest.main()