import unittest from unittest.mock import patch, MagicMock import json import os import pandas as pd import src.codes.code_funcs as code_funcs import src.codes.code_constants as code_constants from src.enums.delimiters import Delimiter class TestCodeFuncs(unittest.TestCase): def test_clean_service(self): """Test the clean_service function with various inputs""" # Test empty input self.assertEqual(code_funcs.clean_service(""), "") self.assertEqual(code_funcs.clean_service(None), "") # Test uppercase conversion self.assertEqual(code_funcs.clean_service("test service"), "TEST SERVICE") # Test synonym mapping with patch.dict(code_constants.SYNONYM_MAP, {"CT": "COMPUTED TOMOGRAPHY"}): self.assertEqual(code_funcs.clean_service("CT SCAN"), "COMPUTED TOMOGRAPHY SCAN") # Test removal list with patch.object(code_constants, "REMOVAL_LIST", ["SCAN"]): self.assertEqual(code_funcs.clean_service("CT SCAN"), "CT") # Test unlisted service self.assertEqual(code_funcs.clean_service("UNLISTED PROCEDURE"), "UNLISTED") # Test stop words with patch.object(code_constants, "STOP_WORD_LIST", ["THE", "AND"]): self.assertEqual(code_funcs.clean_service("THE AND"), "") @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 invoking Claude""" # Setup mock responses mock_invoke_claude.return_value = '{"CPT4_PROC_CD": "12345"}' mock_json_load.return_value = {"CPT4_PROC_CD": "12345"} # Test successful extraction result = code_funcs.code_explicit("TEST SERVICE", "test.pdf") mock_invoke_claude.assert_called_once() self.assertEqual(result, {"CPT4_PROC_CD": "12345"}) @patch('src.utils.llm_utils.invoke_claude') @patch('src.utils.string_utils.extract_text_from_delimiters') def test_code_implicit_special(self, mock_extract, mock_invoke_claude): """Test the code_implicit_special function for special case handling""" # Setup mock responses mock_invoke_claude.return_value = "mock_response" mock_extract.return_value = "Drugs" 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 mock_extract.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, _, _, _ = code_funcs.get_embedding_levels(1, implicit_run_dict) self.assertIn("cpt_level1", embedding_levels) self.assertIn("hcpcs_level1", embedding_levels) self.assertIn("rev_level1", embedding_levels) # Test level 2 with only procedure code enabled implicit_run_dict = {"PROCEDURE_CD": True, "REVENUE_CD": False} embedding_levels, _, _, _ = code_funcs.get_embedding_levels(2, implicit_run_dict) self.assertIn("cpt_level2", embedding_levels) self.assertIn("hcpcs_level2", embedding_levels) self.assertNotIn("rev", embedding_levels) @patch('src.utils.embedding_utils.load_faiss_index') @patch('src.codes.code_constants.model.encode') def test_get_match_list(self, mock_encode, mock_load_faiss): """Test the get_match_list function for retrieving matches from embeddings""" # Setup mock responses mock_encoded_vector = MagicMock() mock_encoded_vector.astype.return_value = "mock_vector" mock_encode.return_value = mock_encoded_vector # Create a mock FAISS index mock_index = MagicMock() # Configure the search method to return appropriate values mock_index.search.return_value = ( # Scores array - 2D array with shape [1, top_k] [[0.9, 0.8]], # Indices array - 2D array with shape [1, top_k] [[0, 1]] ) # Setup mock choices for the returned index choices = ["Description A", "Description B"] # Configure load_faiss_index to return our mock objects mock_load_faiss.return_value = (mock_index, None, choices) # Call the function match_list, highest_similarity = code_funcs.get_match_list("TEST SERVICE", ["test_level"], 2) # Verify results self.assertEqual(match_list, ["Description A", "Description B"]) self.assertEqual(highest_similarity, 0.9) # Verify calls were made correctly mock_encode.assert_called_once_with(["TEST SERVICE"], normalize_embeddings=True) # Make sure load_faiss_index was called for the test_level mock_load_faiss.assert_called_once() # Check that search was called with the encoded vector mock_index.search.assert_called_once_with("mock_vector", 2) @patch('src.codes.code_funcs.get_embedding_levels') @patch('src.codes.code_funcs.get_match_list') @patch('src.utils.llm_utils.invoke_claude') @patch('src.utils.string_utils.universal_json_load') def test_code_implicit_rag(self, mock_json_load, mock_invoke_claude, mock_get_match_list, mock_get_embedding_levels): """Test the code_implicit_rag function for RAG-based code extraction""" # Setup mock responses mock_get_embedding_levels.return_value = (["test_level"], {"12345": "Test Procedure"}, {}, {"678": "Test Revenue"}) mock_get_match_list.return_value = (["Test Procedure"], 0.9) mock_invoke_claude.return_value = '["Test Procedure"]' mock_json_load.return_value = ["Test Procedure"] result = code_funcs.code_implicit_rag("TEST SERVICE", {"PROCEDURE_CD": True, "REVENUE_CD": True}, "test.pdf") self.assertEqual(result["PROCEDURE_CD"], "12345") self.assertEqual(result["PROCEDURE_CD_DESC"], "Test Procedure") self.assertIn("CODE_METHODOLOGY", result) # Test exception handling mock_json_load.side_effect = Exception("Test error") result = code_funcs.code_implicit_rag("TEST SERVICE", {"PROCEDURE_CD": True, "REVENUE_CD": True}, "test.pdf") self.assertIn("CODE_METHODOLOGY", result) @patch('src.utils.llm_utils.invoke_claude') @patch('src.utils.string_utils.extract_text_from_delimiters') def test_code_last_check(self, mock_extract, mock_invoke_claude): """Test the code_last_check function for categorizing non-matched services""" # Setup mock responses mock_invoke_claude.return_value = "mock_response" mock_extract.return_value = "Specific" result = code_funcs.code_last_check("SPECIAL SERVICE", "test.pdf") self.assertEqual(result, "Specific") # Test generic case mock_extract.return_value = "Generic" result = code_funcs.code_last_check("GENERIC SERVICE", "test.pdf") self.assertEqual(result, "Generic") @patch('src.utils.llm_utils.invoke_claude') @patch('src.utils.string_utils.universal_json_load') def test_fill_bill_type(self, mock_json_load, mock_invoke_claude): """Test the fill_bill_type function for populating bill type information""" # Setup mock responses mock_invoke_claude.return_value = '["Inpatient Hospital"]' mock_json_load.return_value = ["Inpatient Hospital"] # Mock bill type mappings with patch.dict(code_constants.bill_type_reverse_mapping, {"Inpatient Hospital": "11X"}): result = code_funcs.fill_bill_type("INPATIENT SERVICE", {}) self.assertEqual(result["BILL_TYPE_CD"], "11X") self.assertEqual(result["BILL_TYPE_CD_DESC"], "Inpatient Hospital") 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 mock responses 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 mock_explicit.return_value = {"PROCEDURE_CD": "12345"} result = code_funcs.extract_codes_from_service({"SERVICE_TERM": "TEST SERVICE"}) 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.assertEqual(result["PROCEDURE_CD_DESC"], "UNLISTED") # Test category match 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.assertEqual(result["CODE_METHODOLOGY"], "Implicit - Letter Category") # Test special case match 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.assertEqual(result["CODE_METHODOLOGY"], "Implicit - Special Case") # Test RAG match mock_implicit_special.return_value = {} mock_implicit_rag.return_value = {"PROCEDURE_CD": "12345", "CODE_METHODOLOGY": "Implicit - Level 1"} result = code_funcs.extract_codes_from_service({"SERVICE_TERM": "TEST SERVICE"}) self.assertEqual(result["CODE_METHODOLOGY"], "Implicit - Level 1") # Test no match mock_implicit_rag.return_value = {} mock_last_check.return_value = "Specific" result = code_funcs.extract_codes_from_service({"SERVICE_TERM": "TEST SERVICE"}) self.assertEqual(result["CODE_METHODOLOGY"], "Specific - No Match") def test_fill_claim_type(self): """Test the fill_claim_type function for determining claim types""" # Test with existing claim types answer_dicts = [ {"AARETE_DERIVED_CLAIM_TYPE_CD": "H"}, {"AARETE_DERIVED_CLAIM_TYPE_CD": "H"}, {"AARETE_DERIVED_CLAIM_TYPE_CD": "M"}, {} ] result = code_funcs.fill_claim_type(answer_dicts) self.assertEqual(result[3]["AARETE_DERIVED_CLAIM_TYPE_CD"], "H") # Test with no claim types answer_dicts = [{}, {}] result = code_funcs.fill_claim_type(answer_dicts) self.assertEqual(result, answer_dicts) @patch('src.codes.code_funcs.fill_claim_type') @patch('src.codes.code_funcs.extract_codes_from_service') def test_code_breakout(self, mock_extract_codes, mock_fill_claim_type): """Test the code_breakout function for processing a DataFrame of service records""" # Create a test DataFrame input_df = pd.DataFrame([ {"SERVICE_TERM": "SERVICE A"}, {"SERVICE_TERM": "SERVICE B"} ]) # Setup mock responses for the records after conversion to list of dicts mock_fill_claim_type.return_value = [ {"SERVICE_TERM": "SERVICE A"}, {"SERVICE_TERM": "SERVICE B"} ] # Setup extract_codes_from_service mock to return augmented records mock_extract_codes.side_effect = [ {"SERVICE_TERM": "SERVICE A", "PROCEDURE_CD": "12345"}, {"SERVICE_TERM": "SERVICE B", "PROCEDURE_CD": "67890"} ] # Call the function result_df = code_funcs.code_breakout(input_df) # Verify the result is a DataFrame self.assertIsInstance(result_df, pd.DataFrame) # Verify the correct procedure codes were added self.assertEqual(result_df.iloc[0]["PROCEDURE_CD"], "12345") self.assertEqual(result_df.iloc[1]["PROCEDURE_CD"], "67890") # Verify the DataFrame has the expected number of rows self.assertEqual(len(result_df), 2) # Verify fill_claim_type was called with the correct data mock_fill_claim_type.assert_called_once() args, _ = mock_fill_claim_type.call_args self.assertIsInstance(args[0], list) # Verify extract_codes_from_service was called for each record self.assertEqual(mock_extract_codes.call_count, 2) if __name__ == '__main__': unittest.main()