Merged in feature/doczy-codes (pull request #619)
Feature/doczy codes * Ignore service and procedure generic terms * UPdate PT/OT/ST codes * Standardize stop-loss * UPdate for multiple identical values * Modify code funcs * Run highest match Implicit Level first * Fix implicit code bug * Add PT individually to special case * change revenue code conditions * Merge branch 'main' into feature/doczy-codes * Merged main into feature/doczy-codes * Update main * Relocate create_code_embeddings back to scripts * Remove investment code_funcs (use codes.code_funcs now) * Add docstrings * stop-loss * Update unit tests * Remove test * Remove reimb_dates * fix dtypes * Remove prints * update prompt * Update code breakout unit test * Merge branch 'main' into feature/doczy-codes * Merge branch 'main' into feature/doczy-codes Approved-by: Alex Galarce
This commit is contained in:
Katon Minhas
@@ -0,0 +1,317 @@
|
||||
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()
|
||||
Reference in New Issue
Block a user