Files
doczyai-pipelines/fieldExtraction/tests/test_postprocess.py
T
Katon Minhas 2b9bf50c2e Merged in feature/effective-dt (pull request #752)
Feature/effective dt

* move prompt_smart_chunked to prompt_calls.py

* refactor pass_smart_chunked_to_full_context

* Mypy fix

* Update field to EFFECTIVE_DT

* move to prompt_calls

* Update vision funcs

* Restructure prompt-based date derivation

* E2E passed

* Remove test

* Merge branch 'main' into feature/effective-dt

* remove deprecated postprocessing_func

* Docstrings

* refactor - send global lesser of to prompt_calls.py

* remove deprecated function

* more docstrings

* Merged main into feature/effective-dt

* Merged main into feature/effective-dt

* added base field for aarete derived effective date


Approved-by: VenkataKrishna Reddy Avula
2025-10-30 15:28:39 +00:00

539 lines
22 KiB
Python

import unittest
from unittest.mock import MagicMock, patch
import pandas as pd
import pytest
from src.investment.postprocessing_funcs import (
deduplicate_provider_columns,
flatten_singleton_string_list, format_rate_fields_with_commas,
normalize_auto_renewal_term, normalize_cpt_fields,
normalize_indicator_field, process_patient_age_range, remove_hyphens,
remove_redundant_reimb_info, rename_columns, validate_and_reformat_date,
fill_empty_dynamic)
class TestPostprocessFunctions(unittest.TestCase):
def test_normalize_indicator_field(self):
self.assertEqual(normalize_indicator_field("Y"), "Y")
self.assertEqual(normalize_indicator_field("y"), "Y")
self.assertEqual(normalize_indicator_field("N"), "N")
self.assertEqual(normalize_indicator_field(""), "N")
self.assertEqual(normalize_indicator_field(None), "N")
def test_format_rate_fields_with_commas(self):
# Test numeric values
self.assertEqual(format_rate_fields_with_commas("1234.567"), "1,234.57")
self.assertEqual(format_rate_fields_with_commas("1000"), "1,000.00")
self.assertEqual(format_rate_fields_with_commas(1234.567), "1,234.57")
self.assertEqual(format_rate_fields_with_commas(1000), "1,000.00")
def test_remove_hyphens(self):
self.assertEqual(remove_hyphens("123-45-6789"), "123456789")
self.assertEqual(remove_hyphens("123-456"), "123456")
self.assertEqual(remove_hyphens(None), "")
self.assertEqual(remove_hyphens(""), "")
def test_flatten_singleton_string_list(self):
self.assertEqual(flatten_singleton_string_list("['123']"), "123")
self.assertEqual(flatten_singleton_string_list("['123', '456']"), "123, 456")
self.assertEqual(flatten_singleton_string_list("11"), "11")
self.assertEqual(flatten_singleton_string_list("invalid"), "invalid")
self.assertEqual(flatten_singleton_string_list(None), "")
def test_rename_columns(self):
"""Tests the rename_columns function to ensure it correctly renames specified columns.
Tests:
1. Basic column renaming from PROCEDURE_CD to CPT4_PROC_CD
2. Multiple columns being renamed
3. Handling of columns not in the rename map
4. Empty DataFrame
"""
# Test case 1: Basic column renaming
input_df1 = pd.DataFrame(
{
"PROCEDURE_CD": ["12345", "67890"],
"PROCEDURE_CD_DESC": ["Test Procedure", "Another Procedure"],
"OTHER_COLUMN": ["value1", "value2"],
}
)
expected_df1 = pd.DataFrame(
{
"CPT4_PROC_CD": ["12345", "67890"],
"CPT4_PROC_CD_DESC": ["Test Procedure", "Another Procedure"],
"OTHER_COLUMN": ["value1", "value2"],
}
)
result_df1 = rename_columns(input_df1)
pd.testing.assert_frame_equal(result_df1, expected_df1)
# Test case 2: Only some columns need renaming
input_df2 = pd.DataFrame(
{"PROCEDURE_CD": ["12345", "67890"], "OTHER_COLUMN": ["value1", "value2"]}
)
expected_df2 = pd.DataFrame(
{"CPT4_PROC_CD": ["12345", "67890"], "OTHER_COLUMN": ["value1", "value2"]}
)
result_df2 = rename_columns(input_df2)
pd.testing.assert_frame_equal(result_df2, expected_df2)
# Test case 3: None of the columns need renaming
input_df3 = pd.DataFrame({"COLUMN_A": ["a", "b"], "COLUMN_B": ["c", "d"]})
result_df3 = rename_columns(input_df3)
pd.testing.assert_frame_equal(result_df3, input_df3) # Should be unchanged
# Test case 4: Empty DataFrame
empty_df = pd.DataFrame()
result_empty_df = rename_columns(empty_df)
pd.testing.assert_frame_equal(result_empty_df, empty_df) # Should be unchanged
def test_normalize_auto_renewal_term(self):
self.assertEqual(normalize_auto_renewal_term("12 months"), "1 year")
self.assertEqual(normalize_auto_renewal_term("one year"), "1 year")
self.assertEqual(normalize_auto_renewal_term("month to month"), "1 month")
self.assertEqual(normalize_auto_renewal_term("(12) 12 months"), "1 year")
self.assertEqual(normalize_auto_renewal_term(None), "")
def test_normalize_cpt_fields(self):
self.assertEqual(normalize_cpt_fields("[123, 456]"), "['123', '456']")
self.assertEqual(normalize_cpt_fields("123-456"), "['123-456']")
self.assertEqual(
normalize_cpt_fields("['T0000-T9999, S0000-S9999']"),
"['T0000-T9999', 'S0000-S9999']",
)
self.assertEqual(normalize_cpt_fields("123"), "['123']")
self.assertEqual(normalize_cpt_fields(None), "")
self.assertEqual(normalize_cpt_fields(123), "['123']")
def test_process_patient_age_range(self):
"""Tests the process_patient_age_range function with various age range formats.
Tests:
1. Standard hyphenated ranges (e.g., "0-18")
2. Single age values
3. Text descriptions with "to"
4. Text descriptions with "and under"
5. Special cases like "newborn"
6. Empty/None values
7. Invalid formats
8. Missing PATIENT_AGE_RANGE column
"""
# Test case 1: DataFrame with PATIENT_AGE_RANGE column
input_df = pd.DataFrame(
{
"PATIENT_AGE_RANGE": [
"0-18", # Standard hyphenated range
"21", # Single age
]
}
)
result_df = process_patient_age_range(input_df)
# Verify columns
self.assertIn("PATIENT_AGE_MIN", result_df.columns)
self.assertIn("PATIENT_AGE_MAX", result_df.columns)
self.assertNotIn("PATIENT_AGE_RANGE", result_df.columns)
# Expected values
expected_min = ["0", "21"]
expected_max = ["18", "21"]
# Check transformations
pd.testing.assert_series_equal(
result_df["PATIENT_AGE_MIN"],
pd.Series(expected_min, name="PATIENT_AGE_MIN"),
check_dtype=False,
)
pd.testing.assert_series_equal(
result_df["PATIENT_AGE_MAX"],
pd.Series(expected_max, name="PATIENT_AGE_MAX"),
check_dtype=False,
)
# Test case 2: DataFrame without PATIENT_AGE_RANGE column
input_df_no_age = pd.DataFrame({"OTHER_COLUMN": ["value1", "value2"]})
result_df_no_age = process_patient_age_range(input_df_no_age)
# Verify the DataFrame is unchanged
pd.testing.assert_frame_equal(input_df_no_age, result_df_no_age)
# Test case 3: Empty DataFrame
empty_df = pd.DataFrame()
result_empty_df = process_patient_age_range(empty_df)
# Verify empty DataFrame is unchanged
pd.testing.assert_frame_equal(empty_df, result_empty_df)
def test_validate_and_reformat_date(self):
"""Tests validate_and_reformat_date with various date formats.
Tests:
1. Date already in YYYY/MM/DD format
2. Common alternative formats (YYYY-MM-DD, MM/DD/YYYY, etc.)
3. Invalid date formats
4. None and non-string values
"""
# Date already in correct format
self.assertEqual(validate_and_reformat_date("2023/01/15"), "2023/01/15")
# Test various date formats that should be reformatted
self.assertEqual(validate_and_reformat_date("2023-01-15"), "2023/01/15")
self.assertEqual(validate_and_reformat_date("01/15/2023"), "2023/01/15")
self.assertEqual(validate_and_reformat_date("15-Jan-2023"), "2023/01/15")
# Test invalid formats - should return the original string
self.assertEqual(validate_and_reformat_date("Invalid date"), "Invalid date")
self.assertEqual(validate_and_reformat_date("01-15"), "01-15")
# Test None and non-string values
self.assertEqual(validate_and_reformat_date(None), None)
self.assertEqual(validate_and_reformat_date(12345), 12345)
def test_remove_redundant_reimb_info(self):
"""Tests remove_redundant_reimb_date function.
Tests:
1. When reimbursement dates match derived dates (should remove)
2. When reimbursement dates differ from derived dates (should keep)
3. When only some rows match (should remove only matching rows)
4. When columns are missing (should return unchanged DataFrame)
"""
# Test case 1: When dates match (should remove)
input_df1 = pd.DataFrame(
{
"REIMB_EFFECTIVE_DT": ["2023/01/01", "2023/02/01"],
"REIMB_TERMINATION_DT": ["2023/12/31", "2023/12/31"],
"AARETE_DERIVED_EFFECTIVE_DT": ["2023/01/01", "2023/02/01"],
"AARETE_DERIVED_TERMINATION_DT": ["2023/12/31", "2023/12/31"],
}
)
expected_df1 = pd.DataFrame(
{
"REIMB_EFFECTIVE_DT": ["", ""],
"REIMB_TERMINATION_DT": ["", ""],
"AARETE_DERIVED_EFFECTIVE_DT": ["2023/01/01", "2023/02/01"],
"AARETE_DERIVED_TERMINATION_DT": ["2023/12/31", "2023/12/31"],
}
)
result_df1 = remove_redundant_reimb_info(input_df1)
pd.testing.assert_frame_equal(result_df1, expected_df1)
# Test case 2: When dates differ (should keep)
input_df2 = pd.DataFrame(
{
"REIMB_EFFECTIVE_DT": ["2023/01/15", "2023/02/15"],
"REIMB_TERMINATION_DT": ["2023/12/15", "2023/12/15"],
"AARETE_DERIVED_EFFECTIVE_DT": ["2023/01/01", "2023/02/01"],
"AARETE_DERIVED_TERMINATION_DT": ["2023/12/31", "2023/12/31"],
}
)
# Result should be unchanged
result_df2 = remove_redundant_reimb_info(input_df2)
pd.testing.assert_frame_equal(result_df2, input_df2)
# Test case 3: Mixed case - some match, some don't
input_df3 = pd.DataFrame(
{
"REIMB_EFFECTIVE_DT": ["2023/01/01", "2023/02/15"],
"REIMB_TERMINATION_DT": ["2023/12/31", "2023/12/15"],
"AARETE_DERIVED_EFFECTIVE_DT": ["2023/01/01", "2023/02/01"],
"AARETE_DERIVED_TERMINATION_DT": ["2023/12/31", "2023/12/31"],
}
)
expected_df3 = pd.DataFrame(
{
"REIMB_EFFECTIVE_DT": ["", "2023/02/15"],
"REIMB_TERMINATION_DT": ["", "2023/12/15"],
"AARETE_DERIVED_EFFECTIVE_DT": ["2023/01/01", "2023/02/01"],
"AARETE_DERIVED_TERMINATION_DT": ["2023/12/31", "2023/12/31"],
}
)
result_df3 = remove_redundant_reimb_info(input_df3)
pd.testing.assert_frame_equal(result_df3, expected_df3)
# Test case 4: Missing columns
input_df4 = pd.DataFrame(
{
"REIMB_EFFECTIVE_DT": ["2023/01/01", "2023/02/01"],
"OTHER_COLUMN": ["value1", "value2"],
}
)
# Result should be unchanged
result_df4 = remove_redundant_reimb_info(input_df4)
pd.testing.assert_frame_equal(result_df4, input_df4)
def test_deduplicate_provider_columns(self):
"""Tests deduplicate_provider_columns function.
Tests:
1. Basic deduplication - removes GROUP values from OTHER fields
2. Internal deduplication - removes duplicates within OTHER fields
3. Combined scenario - both GROUP removal and internal deduplication
4. Empty/missing values handling
5. Missing columns - should return unchanged DataFrame
6. Empty DataFrame
"""
# Test case 1: Basic GROUP removal
input_df1 = pd.DataFrame(
{
"PROV_GROUP_TIN": ["123456789"],
"PROV_GROUP_NPI": ["1234567890"],
"PROV_GROUP_NAME_FULL": ["Main Hospital"],
"PROV_OTHER_TIN": ["123456789|987654321"],
"PROV_OTHER_NPI": ["1234567890|0987654321"],
"PROV_OTHER_NAME_FULL": ["Main Hospital|Other Clinic"],
}
)
expected_df1 = pd.DataFrame(
{
"PROV_GROUP_TIN": ["123456789"],
"PROV_GROUP_NPI": ["1234567890"],
"PROV_GROUP_NAME_FULL": ["Main Hospital"],
"PROV_OTHER_TIN": ["987654321"],
"PROV_OTHER_NPI": ["0987654321"],
"PROV_OTHER_NAME_FULL": ["Other Clinic"],
}
)
result_df1 = deduplicate_provider_columns(input_df1)
pd.testing.assert_frame_equal(result_df1, expected_df1)
# Test case 2: Internal deduplication (your original example)
input_df2 = pd.DataFrame(
{
"PROV_GROUP_TIN": ["061798267"],
"PROV_GROUP_NPI": ["1111111111"],
"PROV_GROUP_NAME_FULL": ["Group Practice"],
"PROV_OTHER_TIN": [
"061798267|061798267|UNKNOWN|UNKNOWN|061992277|061798267"
],
"PROV_OTHER_NPI": ["1111111111|2222222222|2222222222|UNKNOWN"],
"PROV_OTHER_NAME_FULL": [
"Group Practice|Other Practice|Other Practice|UNKNOWN"
],
}
)
expected_df2 = pd.DataFrame(
{
"PROV_GROUP_TIN": ["061798267"],
"PROV_GROUP_NPI": ["1111111111"],
"PROV_GROUP_NAME_FULL": ["Group Practice"],
"PROV_OTHER_TIN": ["061992277"],
"PROV_OTHER_NPI": ["2222222222"],
"PROV_OTHER_NAME_FULL": ["Other Practice"],
}
)
result_df2 = deduplicate_provider_columns(input_df2)
pd.testing.assert_frame_equal(result_df2, expected_df2)
# Test case 3: Empty OTHER fields after deduplication
input_df3 = pd.DataFrame(
{
"PROV_GROUP_TIN": ["123456789"],
"PROV_GROUP_NPI": ["1234567890"],
"PROV_GROUP_NAME_FULL": ["Main Hospital"],
"PROV_OTHER_TIN": ["123456789|123456789|UNKNOWN"],
"PROV_OTHER_NPI": ["1234567890|UNKNOWN|UNKNOWN"],
"PROV_OTHER_NAME_FULL": ["Main Hospital|UNKNOWN"],
}
)
expected_df3 = pd.DataFrame(
{
"PROV_GROUP_TIN": ["123456789"],
"PROV_GROUP_NPI": ["1234567890"],
"PROV_GROUP_NAME_FULL": ["Main Hospital"],
"PROV_OTHER_TIN": [""],
"PROV_OTHER_NPI": [""],
"PROV_OTHER_NAME_FULL": [""],
}
)
result_df3 = deduplicate_provider_columns(input_df3)
pd.testing.assert_frame_equal(result_df3, expected_df3)
# Test case 4: Multiple rows
input_df4 = pd.DataFrame(
{
"PROV_GROUP_TIN": ["111111111", "222222222"],
"PROV_GROUP_NPI": ["1111111111", "2222222222"],
"PROV_GROUP_NAME_FULL": ["Hospital A", "Hospital B"],
"PROV_OTHER_TIN": [
"111111111|333333333",
"444444444|222222222|444444444",
],
"PROV_OTHER_NPI": ["3333333333|1111111111", "4444444444|2222222222"],
"PROV_OTHER_NAME_FULL": ["Clinic C|Hospital A", "Clinic D|Hospital B"],
}
)
expected_df4 = pd.DataFrame(
{
"PROV_GROUP_TIN": ["111111111", "222222222"],
"PROV_GROUP_NPI": ["1111111111", "2222222222"],
"PROV_GROUP_NAME_FULL": ["Hospital A", "Hospital B"],
"PROV_OTHER_TIN": ["333333333", "444444444"],
"PROV_OTHER_NPI": ["3333333333", "4444444444"],
"PROV_OTHER_NAME_FULL": ["Clinic C", "Clinic D"],
}
)
result_df4 = deduplicate_provider_columns(input_df4)
pd.testing.assert_frame_equal(result_df4, expected_df4)
# Test case 5: Missing columns - should return unchanged
input_df5 = pd.DataFrame(
{"PROV_GROUP_TIN": ["123456789"], "OTHER_COLUMN": ["value1"]}
)
result_df5 = deduplicate_provider_columns(input_df5)
pd.testing.assert_frame_equal(result_df5, input_df5)
# Test case 6: Empty DataFrame
empty_df = pd.DataFrame()
result_empty_df = deduplicate_provider_columns(empty_df)
pd.testing.assert_frame_equal(empty_df, result_empty_df)
# Test case 7: Empty OTHER fields (already empty strings)
input_df7 = pd.DataFrame(
{
"PROV_GROUP_TIN": ["123456789"],
"PROV_GROUP_NPI": ["1234567890"],
"PROV_GROUP_NAME_FULL": ["Main Hospital"],
"PROV_OTHER_TIN": [""],
"PROV_OTHER_NPI": [""],
"PROV_OTHER_NAME_FULL": [""],
}
)
result_df7 = deduplicate_provider_columns(input_df7)
pd.testing.assert_frame_equal(result_df7, input_df7)
def test_fill_empty_dynamic(self):
"""Tests the fill_empty_dynamic function that fills NA values with common values from the same EXHIBIT_PAGE group.
Tests:
1. Basic filling - fills NA values with the common value for the same EXHIBIT_PAGE
2. Multiple columns - correctly fills multiple columns independently
3. Multiple file/page combinations - respects FILE_NAME and EXHIBIT_PAGE boundaries
4. No common value - doesn't fill when multiple non-NA values exist
5. All NA values - doesn't fill when all values are NA
6. Missing columns - returns unchanged DataFrame when key columns are missing
"""
# Test case 1: Basic filling for a single column
input_df1 = pd.DataFrame({
"FILE_NAME": ["file1.pdf", "file1.pdf", "file1.pdf"],
"EXHIBIT_PAGE": ["1.0", "1.0", "1.0"],
"AARETE_DERIVED_LOB": ["Commercial", "", None]
})
expected_df1 = pd.DataFrame({
"FILE_NAME": ["file1.pdf", "file1.pdf", "file1.pdf"],
"EXHIBIT_PAGE": ["1.0", "1.0", "1.0"],
"AARETE_DERIVED_LOB": ["Commercial", "Commercial", "Commercial"]
})
result_df1 = fill_empty_dynamic(input_df1)
pd.testing.assert_frame_equal(result_df1, expected_df1)
# Test case 2: Multiple columns
input_df2 = pd.DataFrame({
"FILE_NAME": ["file1.pdf", "file1.pdf", "file1.pdf"],
"EXHIBIT_PAGE": ["1.0", "1.0", "1.0"],
"AARETE_DERIVED_LOB": ["Commercial", "", None],
"AARETE_DERIVED_PRODUCT": ["Product A", None, ""]
})
expected_df2 = pd.DataFrame({
"FILE_NAME": ["file1.pdf", "file1.pdf", "file1.pdf"],
"EXHIBIT_PAGE": ["1.0", "1.0", "1.0"],
"AARETE_DERIVED_LOB": ["Commercial", "Commercial", "Commercial"],
"AARETE_DERIVED_PRODUCT": ["Product A", "Product A", "Product A"]
})
result_df2 = fill_empty_dynamic(input_df2)
pd.testing.assert_frame_equal(result_df2, expected_df2)
# Test case 3: Multiple file/page combinations
input_df3 = pd.DataFrame({
"FILE_NAME": ["file1.pdf", "file1.pdf", "file1.pdf", "file1.pdf", "file2.pdf", "file2.pdf"],
"EXHIBIT_PAGE": ["1.0", "1.0", "2.0", "2.0", "1.0", "1.0"],
"AARETE_DERIVED_LOB": ["Commercial", "", "Medicare", None, "Medicaid", ""]
})
expected_df3 = pd.DataFrame({
"FILE_NAME": ["file1.pdf", "file1.pdf", "file1.pdf", "file1.pdf", "file2.pdf", "file2.pdf"],
"EXHIBIT_PAGE": ["1.0", "1.0", "2.0", "2.0", "1.0", "1.0"],
"AARETE_DERIVED_LOB": ["Commercial", "Commercial", "Medicare", "Medicare", "Medicaid", "Medicaid"]
})
result_df3 = fill_empty_dynamic(input_df3)
pd.testing.assert_frame_equal(result_df3, expected_df3)
# Test case 4: No common value (multiple non-NA values exist)
input_df4 = pd.DataFrame({
"FILE_NAME": ["file1.pdf", "file1.pdf", "file1.pdf"],
"EXHIBIT_PAGE": ["1.0", "1.0", "1.0"],
"AARETE_DERIVED_LOB": ["Commercial", "Medicare", None]
})
expected_df4 = pd.DataFrame({
"FILE_NAME": ["file1.pdf", "file1.pdf", "file1.pdf"],
"EXHIBIT_PAGE": ["1.0", "1.0", "1.0"],
"AARETE_DERIVED_LOB": ["Commercial", "Medicare", None] # Should remain unchanged
})
result_df4 = fill_empty_dynamic(input_df4)
pd.testing.assert_frame_equal(result_df4, expected_df4)
# Test case 5: All NA values
input_df5 = pd.DataFrame({
"FILE_NAME": ["file1.pdf", "file1.pdf", "file1.pdf"],
"EXHIBIT_PAGE": ["1.0", "1.0", "1.0"],
"AARETE_DERIVED_LOB": [None, "", None]
})
expected_df5 = pd.DataFrame({
"FILE_NAME": ["file1.pdf", "file1.pdf", "file1.pdf"],
"EXHIBIT_PAGE": ["1.0", "1.0", "1.0"],
"AARETE_DERIVED_LOB": [None, "", None] # Should remain unchanged
})
result_df5 = fill_empty_dynamic(input_df5)
pd.testing.assert_frame_equal(result_df5, expected_df5)
# Test case 6: Missing EXHIBIT_PAGE column
input_df6 = pd.DataFrame({
"FILE_NAME": ["file1.pdf", "file1.pdf"],
"AARETE_DERIVED_LOB": ["Commercial", None]
})
result_df6 = fill_empty_dynamic(input_df6)
pd.testing.assert_frame_equal(result_df6, input_df6) # Should remain unchanged
# Test case 7: Empty DataFrame
empty_df = pd.DataFrame()
result_empty = fill_empty_dynamic(empty_df)
pd.testing.assert_frame_equal(result_empty, empty_df) # Should remain unchanged
if __name__ == "__main__":
unittest.main()