jacob-mathison/doczyai-pipelines
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Activity
Files
2da639e59b756c30bfaf22f4862f636a95d80829
doczyai-pipelines/src/tests/test_postprocess.py
T
Venkatakrishna Reddy Avula 2da639e59b Merged in feature/DAIP2-2314-DAIP2-1687-hybrid (pull request #993) 
Feature/DAIP2-2314 DAIP2 1687 hybrid

* remove -files from s3 prefix requirements

* Resolve input paths

* fix: VendorProcessor.process_file returns (df, None) tuple

runner.safe_process_file unpacks the result as (cc_df, dashboard_df), so
returning a single DataFrame caused every vendor/generic file to fail with
"too many values to unpack (expected 2)" — Python iterates DataFrame columns
during unpacking. Vendor pipelines have no dashboard variant; second slot is
None and the existing `dashboard_result is not None` guard in runner.py
already handles it.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

* DAIP2-2314 + DAIP2-1687: pad DYNAMIC_PRIMARY + DYNAMIC_PRIMARY_ENTITY_CLASSIFICATION over 1024-token cache floor

- Pad DYNAMIC_PRIMARY_INSTRUCTION with three new sections: [SCOPE BOUNDARIES], [SOURCE TEXT INTERPRETATION], [REASONING DISCIPLINE], plus a [WORKED EXAMPLES] block. Estimated tokens: 447 -> 1117 (Sonnet 4.5 1024-min, +93 margin). All additions reinforce existing rules (sibling-field separation, alias mapping, pricing-vs-LOB distinction, contrastive-clause exclusion, exhibit-header binding) — no new directives that could bias extraction.
- Pad DYNAMIC_PRIMARY_ENTITY_CLASSIFICATION_INSTRUCTION with a [FINAL CHECKLIST BEFORE OUTPUT] block. Estimated tokens: 956 -> 1101 (Sonnet 4.5 1024-min, +77 margin). Reinforces the existing 4-step anti-duplication protocol and JSON shape requirements.
- Register both new entries in cache_registry: DYNAMIC_PRIMARY_ENTITY_CLASSIFICATION as INSTRUCTION_PLUS_CONTEXT (caches at warm-up), DYNAMIC_PRIMARY_ENTITIES as CONTEXT (instruction is intentionally short; CONTEXT c…
* black format fix

* Merged dev into feature/DAIP2-2314-DAIP2-1687-hybrid

* fixed raw lob values in base lob field mapping and composite entities fix

* black format fix

* fixed LOB Program output issues

* issue fixes

* remove debugging code

* Updated prompts

* updated additional instructions

* Update Program-->LOB

* LLM-based AD-Program/Product mapping to LOB even when there is a crosswalk

* black format fix

* Merged dev into feature/DAIP2-2314-DAIP2-1687-hybrid

* added logging in prompt call tracking

* added updated logging in prompt call tracking

* aaded min cache token per usage label

* added cache registry for dynamic primary mapping prompt calls

* reolved mapping prompts ambiguities

* black format fix

* Phase 2 modifications added

* reverted phase 2 modifications


Approved-by: Katon Minhas
2026-05-08 15:52:32 +00:00

1338 lines
53 KiB
Python
Raw Blame History

import json
import os
import tempfile
import unittest
from unittest.mock import MagicMock, patch
import pandas as pd
import pytest
from src.pipelines.shared.postprocessing.postprocessing_funcs import (
blank_uom_for_default_flat_rate_cc,
clean_na_values,
deduplicate_provider_columns,
fill_claim_type_from_title,
flatten_singleton_string_list,
format_as_json_list,
format_rate_fields_with_commas,
normalize_auto_renewal_term,
normalize_cpt_fields,
normalize_indicator_field,
process_patient_age_range,
remove_redundant_reimb_info,
rename_columns,
standardize_reimb_method_and_fee_schedule,
validate_and_reformat_date,
)
from src.pipelines.shared.postprocessing import aarete_derived
from src.pipelines.shared.postprocessing.postprocess import standard_postprocess
from src.constants.constants import Constants
class TestPostprocessFunctions(unittest.TestCase):
def test_normalize_indicator_field(self):
"""Test normalize_indicator_field with various inputs including blanks/None/NaN."""
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")
# Test with whitespace-only strings
self.assertEqual(normalize_indicator_field(" "), "N")
self.assertEqual(normalize_indicator_field("\t\n"), "N")
# Test with NaN (float)
import numpy as np
self.assertEqual(normalize_indicator_field(np.nan), "N")
self.assertEqual(normalize_indicator_field(float("nan")), "N")
# Test with pd.NA
self.assertEqual(normalize_indicator_field(pd.NA), "N")
# Test that non-Y values become N
self.assertEqual(normalize_indicator_field("X"), "N")
self.assertEqual(normalize_indicator_field("Yes"), "N")
self.assertEqual(normalize_indicator_field("1"), "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_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_format_as_json_list_returns_blank_for_empty(self):
"""Test format_as_json_list returns blank instead of [] for empty lists."""
self.assertEqual(format_as_json_list([]), "")
self.assertEqual(format_as_json_list("[]"), "")
self.assertEqual(format_as_json_list(None), "")
self.assertEqual(format_as_json_list(""), "")
self.assertEqual(format_as_json_list([""]), "")
self.assertEqual(format_as_json_list(["[]"]), "")
self.assertEqual(format_as_json_list(["Valid"]), '["Valid"]')
self.assertEqual(format_as_json_list(["A", "B"]), '["A", "B"]')
def test_format_as_json_list_pipe_delimited(self):
"""Test format_as_json_list converts pipe-delimited strings to JSON list."""
self.assertEqual(format_as_json_list("CHIP|MMC"), '["CHIP", "MMC"]')
self.assertEqual(format_as_json_list("A|B|C"), '["A", "B", "C"]')
self.assertEqual(format_as_json_list(" Single "), '["Single"]')
def test_format_as_json_list_clean_special_keep_apostrophe_hyphen(self):
"""Test format_as_json_list keeps letters, digits, comma, space, apostrophe, hyphen only."""
self.assertEqual(
format_as_json_list("Medicare-Medicaid (MM)"),
'["Medicare-Medicaid MM"]',
)
self.assertEqual(
format_as_json_list("Medicare-Medicaid Plan (MMP)"),
'["Medicare-Medicaid Plan MMP"]',
)
self.assertEqual(
format_as_json_list("Children's / Medicaid-Medicare (MM)"),
'["Children\'s Medicaid-Medicare MM"]',
)
self.assertEqual(
format_as_json_list("Plan (MMP), Other (X)"),
'["Plan MMP", "Other X"]',
)
def test_format_as_json_list_strips_wrapping_single_quotes(self):
"""Test format_as_json_list turns [\"'NV'\"] into [\"NV\"] for state fields."""
self.assertEqual(format_as_json_list(["'NV'"]), '["NV"]')
self.assertEqual(format_as_json_list("[\"'NV'\"]"), '["NV"]')
self.assertEqual(format_as_json_list(["'NV'", "'CA'"]), '["NV", "CA"]')
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")
# Pipe-wrapped date (e.g. from extraction) should be stripped and reformatted
self.assertEqual(validate_and_reformat_date("|01/01/2022|"), "2022/01/01")
self.assertEqual(validate_and_reformat_date("|2023/06/15|"), "2023/06/15")
# Datetime with time component should output date-only YYYY/MM/DD
self.assertEqual(
validate_and_reformat_date("2022-01-01 00:00:00"), "2022/01/01"
)
self.assertEqual(
validate_and_reformat_date("2023-06-15 12:30:00.123456"), "2023/06/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
# deduplicate_provider_columns now works with lists, not pipe-delimited strings
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)
# deduplicate_provider_columns now works with lists, not pipe-delimited strings
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
# deduplicate_provider_columns now works with lists, not pipe-delimited strings
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
# deduplicate_provider_columns now works with lists, not pipe-delimited strings
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 lists)
# deduplicate_provider_columns now works with lists, not pipe-delimited 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_claim_type_from_title_mode_fill(self):
"""Test filling empty values using mode from same file."""
input_df = pd.DataFrame(
{
"FILE_NAME": ["file1.pdf", "file1.pdf", "file1.pdf", "file2.pdf"],
"AARETE_DERIVED_CLAIM_TYPE_CD": ["M", "M", "", "H"],
"CONTRACT_TITLE": ["Test", "Test", "Test", "Hospital Agreement"],
}
)
result_df = fill_claim_type_from_title(input_df)
# Empty value in file1 should be filled with "M" (mode)
assert result_df.loc[2, "AARETE_DERIVED_CLAIM_TYPE_CD"] == "M"
# file2 value should remain unchanged
assert result_df.loc[3, "AARETE_DERIVED_CLAIM_TYPE_CD"] == "H"
def test_fill_claim_type_from_title_list_values(self):
"""Test handling list values (bug fix scenario)."""
input_df = pd.DataFrame(
{
"FILE_NAME": ["file1.pdf", "file1.pdf", "file1.pdf"],
"AARETE_DERIVED_CLAIM_TYPE_CD": [["M"], ["M"], ""],
"CONTRACT_TITLE": ["Test", "Test", "Test"],
}
)
result_df = fill_claim_type_from_title(input_df)
# Should handle list values and fill empty row with mode
assert result_df.loc[2, "AARETE_DERIVED_CLAIM_TYPE_CD"] == "M"
def test_fill_claim_type_from_title_professional_keywords(self):
"""Test inferring M from professional keywords in CONTRACT_TITLE."""
input_df = pd.DataFrame(
{
"CONTRACT_TITLE": [
"Physician Services Agreement",
"Professional Provider Agreement",
"Medical Group Contract",
"Participating Provider Agreement",
],
"AARETE_DERIVED_CLAIM_TYPE_CD": ["", "", "", ""],
}
)
result_df = fill_claim_type_from_title(input_df)
# All should be inferred as "M"
assert all(result_df["AARETE_DERIVED_CLAIM_TYPE_CD"] == "M")
def test_fill_claim_type_from_title_ancillary_keywords(self):
"""Test inferring M from ancillary keywords in CONTRACT_TITLE."""
input_df = pd.DataFrame(
{
"CONTRACT_TITLE": [
"Home Health Services Agreement",
"DME Provider Contract",
"Laboratory Services Agreement",
"Behavioral Health Agreement",
"Ambulatory Surgical Center",
],
"AARETE_DERIVED_CLAIM_TYPE_CD": ["", "", "", "", ""],
}
)
result_df = fill_claim_type_from_title(input_df)
# All should be inferred as "M"
assert all(result_df["AARETE_DERIVED_CLAIM_TYPE_CD"] == "M")
def test_fill_claim_type_from_title_institutional_keywords(self):
"""Test inferring H from institutional keywords in CONTRACT_TITLE."""
input_df = pd.DataFrame(
{
"CONTRACT_TITLE": [
"Hospital Services Agreement",
"Institutional Provider Contract",
"Facility Agreement",
"Inpatient Services",
],
"AARETE_DERIVED_CLAIM_TYPE_CD": ["", "", "", ""],
}
)
result_df = fill_claim_type_from_title(input_df)
# All should be inferred as "H"
assert all(result_df["AARETE_DERIVED_CLAIM_TYPE_CD"] == "H")
def test_fill_claim_type_from_title_case_insensitive(self):
"""Test that keyword matching is case-insensitive."""
input_df = pd.DataFrame(
{
"CONTRACT_TITLE": [
"physician services agreement",
"HOSPITAL AGREEMENT",
"AnCiLlArY sErViCeS",
],
"AARETE_DERIVED_CLAIM_TYPE_CD": ["", "", ""],
}
)
result_df = fill_claim_type_from_title(input_df)
assert result_df.loc[0, "AARETE_DERIVED_CLAIM_TYPE_CD"] == "M"
assert result_df.loc[1, "AARETE_DERIVED_CLAIM_TYPE_CD"] == "H"
assert result_df.loc[2, "AARETE_DERIVED_CLAIM_TYPE_CD"] == "M"
def test_fill_claim_type_from_title_preserve_existing(self):
"""Test that existing non-empty values are preserved."""
input_df = pd.DataFrame(
{
"CONTRACT_TITLE": ["Hospital Agreement", "Physician Agreement"],
"AARETE_DERIVED_CLAIM_TYPE_CD": [
"M",
"H",
], # Opposite of what keywords would infer
}
)
result_df = fill_claim_type_from_title(input_df)
# Should preserve existing values even if they don't match keywords
assert result_df.loc[0, "AARETE_DERIVED_CLAIM_TYPE_CD"] == "M"
assert result_df.loc[1, "AARETE_DERIVED_CLAIM_TYPE_CD"] == "H"
def test_fill_claim_type_from_title_no_match(self):
"""Test that no value is set when no keywords match."""
input_df = pd.DataFrame(
{
"CONTRACT_TITLE": ["Generic Contract", "Some Agreement"],
"AARETE_DERIVED_CLAIM_TYPE_CD": ["", ""],
}
)
result_df = fill_claim_type_from_title(input_df)
# Should remain empty when no keywords match
assert result_df.loc[0, "AARETE_DERIVED_CLAIM_TYPE_CD"] == ""
assert result_df.loc[1, "AARETE_DERIVED_CLAIM_TYPE_CD"] == ""
def test_fill_claim_type_from_title_missing_column(self):
"""Test that function handles missing AARETE_DERIVED_CLAIM_TYPE_CD column."""
input_df = pd.DataFrame(
{
"CONTRACT_TITLE": ["Physician Agreement"],
}
)
result_df = fill_claim_type_from_title(input_df)
# Should create the column and infer value
assert "AARETE_DERIVED_CLAIM_TYPE_CD" in result_df.columns
assert result_df.loc[0, "AARETE_DERIVED_CLAIM_TYPE_CD"] == "M"
def test_fill_claim_type_from_title_missing_contract_title(self):
"""Test that function handles missing CONTRACT_TITLE column."""
input_df = pd.DataFrame(
{
"FILE_NAME": ["file1.pdf"],
"AARETE_DERIVED_CLAIM_TYPE_CD": [""],
}
)
result_df = fill_claim_type_from_title(input_df)
# Should return df unchanged (can't infer without CONTRACT_TITLE)
assert result_df.loc[0, "AARETE_DERIVED_CLAIM_TYPE_CD"] == ""
def test_fill_claim_type_from_title_mixed_scenario(self):
"""Test combined mode filling and keyword inference."""
input_df = pd.DataFrame(
{
"FILE_NAME": ["file1.pdf", "file1.pdf", "file2.pdf", "file2.pdf"],
"CONTRACT_TITLE": [
"Test",
"Test",
"Physician Agreement",
"Generic",
],
"AARETE_DERIVED_CLAIM_TYPE_CD": ["H", "", "", ""],
}
)
result_df = fill_claim_type_from_title(input_df)
# Row 1 should be filled with "H" (mode from file1)
assert result_df.loc[1, "AARETE_DERIVED_CLAIM_TYPE_CD"] == "H"
# Row 2 should be inferred as "M" from keyword
assert result_df.loc[2, "AARETE_DERIVED_CLAIM_TYPE_CD"] == "M"
# Row 3 should remain empty (no mode in file2, no keyword match)
assert result_df.loc[3, "AARETE_DERIVED_CLAIM_TYPE_CD"] == ""
def test_fill_claim_type_from_title_empty_list_values(self):
"""Test handling empty list values."""
input_df = pd.DataFrame(
{
"FILE_NAME": ["file1.pdf", "file1.pdf"],
"AARETE_DERIVED_CLAIM_TYPE_CD": [[], "M"],
"CONTRACT_TITLE": ["Physician Agreement", "Test"],
}
)
result_df = fill_claim_type_from_title(input_df)
# Empty list should be treated as empty and filled with mode or inferred
assert result_df.loc[0, "AARETE_DERIVED_CLAIM_TYPE_CD"] == "M"
def test_clean_na_values_string_placeholders(self):
"""Test clean_na_values removes placeholder strings when they're the only value."""
# Test with string placeholders
input_df = pd.DataFrame(
{
"FIELD1": ["N/A", "Valid Value", "UNKNOWN", "NA"],
"FIELD2": ["Valid", "N/A", "Another Valid", "null"],
}
)
result_df = clean_na_values(input_df)
# N/A, UNKNOWN, NA, null should be removed (become empty strings)
self.assertEqual(result_df.loc[0, "FIELD1"], "")
self.assertEqual(result_df.loc[1, "FIELD1"], "Valid Value")
self.assertEqual(result_df.loc[2, "FIELD1"], "")
self.assertEqual(result_df.loc[3, "FIELD1"], "")
self.assertEqual(result_df.loc[0, "FIELD2"], "Valid")
self.assertEqual(result_df.loc[1, "FIELD2"], "")
self.assertEqual(result_df.loc[2, "FIELD2"], "Another Valid")
self.assertEqual(result_df.loc[3, "FIELD2"], "")
def test_clean_na_values_json_list_placeholders(self):
"""Test clean_na_values handles JSON list placeholders."""
# Test with JSON string lists
input_df = pd.DataFrame(
{
"FIELD1": ['["N/A"]', '["Medicare", "N/A"]', '["N/A", "N/A"]'],
"FIELD2": ['["Valid"]', '["N/A"]', '["Medicare"]'],
}
)
result_df = clean_na_values(input_df)
# '["N/A"]' should be removed (only placeholder)
self.assertEqual(result_df.loc[0, "FIELD1"], "")
# '["Medicare", "N/A"]' should be kept (has non-placeholder)
self.assertEqual(result_df.loc[1, "FIELD1"], '["Medicare", "N/A"]')
# '["N/A", "N/A"]' should be removed (all placeholders)
self.assertEqual(result_df.loc[2, "FIELD1"], "")
self.assertEqual(result_df.loc[0, "FIELD2"], '["Valid"]')
self.assertEqual(result_df.loc[1, "FIELD2"], "")
self.assertEqual(result_df.loc[2, "FIELD2"], '["Medicare"]')
def test_clean_na_values_empty_list_representations(self):
"""Test clean_na_values blanks [] and ['[]'] instead of showing them."""
input_df = pd.DataFrame(
{
"FIELD1": ["[]", "['[]']", '["[]"]', "Valid"],
"FIELD2": [["[]"], ["Valid"], [], "[]"],
}
)
result_df = clean_na_values(input_df)
self.assertEqual(result_df.loc[0, "FIELD1"], "")
self.assertEqual(result_df.loc[1, "FIELD1"], "")
self.assertEqual(result_df.loc[2, "FIELD1"], "")
self.assertEqual(result_df.loc[3, "FIELD1"], "Valid")
self.assertEqual(result_df.loc[0, "FIELD2"], "")
self.assertEqual(result_df.loc[1, "FIELD2"], ["Valid"])
self.assertEqual(result_df.loc[2, "FIELD2"], "")
self.assertEqual(result_df.loc[3, "FIELD2"], "")
def test_clean_na_values_python_lists(self):
"""Test clean_na_values handles Python list types."""
# Test with Python lists
input_df = pd.DataFrame(
{
"FIELD1": [["N/A"], ["Medicare", "N/A"], ["N/A", "N/A"]],
"FIELD2": [["Valid"], ["N/A"], ["Medicare"]],
}
)
result_df = clean_na_values(input_df)
# ["N/A"] should be removed
self.assertEqual(result_df.loc[0, "FIELD1"], "")
# ["Medicare", "N/A"] should be kept
self.assertEqual(result_df.loc[1, "FIELD1"], ["Medicare", "N/A"])
# ["N/A", "N/A"] should be removed
self.assertEqual(result_df.loc[2, "FIELD1"], "")
self.assertEqual(result_df.loc[0, "FIELD2"], ["Valid"])
self.assertEqual(result_df.loc[1, "FIELD2"], "")
self.assertEqual(result_df.loc[2, "FIELD2"], ["Medicare"])
def test_clean_na_values_none_and_nan(self):
"""Test clean_na_values handles None and NaN values."""
import numpy as np
input_df = pd.DataFrame(
{
"FIELD1": [None, np.nan, pd.NA, "Valid"],
"FIELD2": ["", " ", "Valid", None],
}
)
result_df = clean_na_values(input_df)
# None, NaN, pd.NA, empty strings should be removed (converted to empty strings)
self.assertEqual(result_df.loc[0, "FIELD1"], "")
self.assertEqual(result_df.loc[1, "FIELD1"], "")
# pd.NA should be converted to empty string
result_val = result_df.loc[2, "FIELD1"]
self.assertTrue(
result_val == "" or pd.isna(result_val),
f"Expected empty string or NA, got: {result_val}",
)
self.assertEqual(result_df.loc[3, "FIELD1"], "Valid")
self.assertEqual(result_df.loc[0, "FIELD2"], "")
self.assertEqual(result_df.loc[1, "FIELD2"], "")
self.assertEqual(result_df.loc[2, "FIELD2"], "Valid")
self.assertEqual(result_df.loc[3, "FIELD2"], "")
def test_clean_na_values_empty_dataframe(self):
"""Test clean_na_values handles empty DataFrame."""
empty_df = pd.DataFrame()
result_df = clean_na_values(empty_df)
pd.testing.assert_frame_equal(result_df, empty_df)
def test_clean_na_values_case_insensitive(self):
"""Test clean_na_values is case-insensitive for placeholder matching."""
input_df = pd.DataFrame(
{
"FIELD1": ["n/a", "N/A", "Na", "unknown", "UNKNOWN", "None", "NULL"],
}
)
result_df = clean_na_values(input_df)
# All should be removed (case-insensitive matching)
for idx in range(len(result_df)):
self.assertEqual(result_df.loc[idx, "FIELD1"], "")
def test_standard_postprocess_cleans_na_and_normalizes_ind(self):
"""Test standard_postprocess removes N/A values and normalizes _IND fields."""
constants = Constants()
input_df = pd.DataFrame(
{
"FILE_NAME": ["test1.pdf", "test2.pdf"],
"PAYER_NAME": [
"N/A",
"Valid Value",
], # Use actual column from FIELD_FORMAT_MAPPING
"AUTO_RENEWAL_IND": ["Y", ""],
"DSH_IND": [None, "N"],
"IME_IND": [" ", "y"],
"NETWORK": ['["N/A"]', '["Medicare"]'],
}
)
result_df = standard_postprocess(input_df, constants)
# N/A should be removed (PAYER_NAME is in FIELD_FORMAT_MAPPING so it will be preserved)
self.assertEqual(result_df.loc[0, "PAYER_NAME"], "")
self.assertEqual(result_df.loc[1, "PAYER_NAME"], "Valid Value")
# _IND fields should be normalized to Y or N
self.assertEqual(result_df.loc[0, "AUTO_RENEWAL_IND"], "Y")
self.assertEqual(result_df.loc[1, "AUTO_RENEWAL_IND"], "N") # Empty becomes N
self.assertEqual(result_df.loc[0, "DSH_IND"], "N") # None becomes N
self.assertEqual(result_df.loc[1, "DSH_IND"], "N")
self.assertEqual(result_df.loc[0, "IME_IND"], "N") # Whitespace becomes N
self.assertEqual(result_df.loc[1, "IME_IND"], "Y") # "y" becomes "Y"
# JSON list with only N/A should be removed
self.assertEqual(result_df.loc[0, "NETWORK"], "")
# JSON list with valid value should be kept
self.assertEqual(result_df.loc[1, "NETWORK"], '["Medicare"]')
def test_standard_postprocess_empty_dataframe(self):
"""Test standard_postprocess handles empty DataFrame."""
constants = Constants()
empty_df = pd.DataFrame()
result_df = standard_postprocess(empty_df, constants)
pd.testing.assert_frame_equal(result_df, empty_df)
def test_standard_postprocess_cleans_na_without_ind_fields(self):
"""Test standard_postprocess cleans N/A values even when no _IND fields are present."""
constants = Constants()
input_df = pd.DataFrame(
{
"FILE_NAME": ["test1.pdf", "test2.pdf"],
"PAYER_NAME": [
"N/A",
"Valid",
], # Use actual column from FIELD_FORMAT_MAPPING
"CONTRACT_TITLE": [
"UNKNOWN",
"Another Valid",
], # Use actual column from FIELD_FORMAT_MAPPING
}
)
result_df = standard_postprocess(input_df, constants)
# Should still clean N/A values (these columns are in FIELD_FORMAT_MAPPING so they will be preserved)
self.assertEqual(result_df.loc[0, "PAYER_NAME"], "")
self.assertEqual(result_df.loc[1, "PAYER_NAME"], "Valid")
self.assertEqual(result_df.loc[0, "CONTRACT_TITLE"], "")
self.assertEqual(result_df.loc[1, "CONTRACT_TITLE"], "Another Valid")
@patch(
"src.pipelines.shared.postprocessing.aarete_derived.prompt_calls.prompt_product_to_lob",
return_value=["Commercial"],
)
def test_fill_na_mapping_derives_lob_from_aarete_derived_product(
self, mock_product_lob
):
"""fill_na_mapping should derive LOB from AARETE_DERIVED_PRODUCT via LLM."""
rows = pd.DataFrame(
[
{
"FILE_NAME": "test_file",
"PRODUCT": ["Connexus", "Synergy"],
"AARETE_DERIVED_PRODUCT": ["Connexus", "Synergy"],
"AARETE_DERIVED_LOB": "",
"LOB": "",
"AARETE_DERIVED_PROGRAM": "",
"LOB_PROGRAM_RELATIONSHIP": "",
"LOB_PRODUCT_RELATIONSHIP": "",
}
]
)
constants = MagicMock()
constants.CROSSWALK_PROGRAM_LOB.mapping = {}
constants.CROSSWALK_PRODUCT_LOB.mapping = {}
constants.CROSSWALK_LOB.mapping = {}
result = aarete_derived.fill_na_mapping(rows, constants)
self.assertEqual(result.iloc[0]["AARETE_DERIVED_LOB"], ["Commercial"])
class TestStandardizeReimbMethodAndFeeScheduleUOM(unittest.TestCase):
"""Tests for UNIT_OF_MEASURE post-processing when DEFAULT_IND='Y' and flat rate."""
def test_cc_default_flat_rate_uom_blank(self):
"""CC only: DEFAULT_IND='Y' and Flat Rate -> UNIT_OF_MEASURE blank."""
df = pd.DataFrame(
{
"DEFAULT_IND": ["Y"],
"AARETE_DERIVED_REIMB_METHOD": ["Flat Rate"],
"UNIT_OF_MEASURE": ["Per Unit"],
}
)
result = blank_uom_for_default_flat_rate_cc(df)
self.assertEqual(result.loc[0, "UNIT_OF_MEASURE"], "")
def test_cc_default_flat_rate_case_insensitive(self):
"""CC only: DEFAULT_IND='Y' and 'flat rate' (lowercase) -> UNIT_OF_MEASURE blank."""
df = pd.DataFrame(
{
"DEFAULT_IND": ["Y"],
"AARETE_DERIVED_REIMB_METHOD": ["flat rate"],
"UNIT_OF_MEASURE": ["Per Visit"],
}
)
result = blank_uom_for_default_flat_rate_cc(df)
self.assertEqual(result.loc[0, "UNIT_OF_MEASURE"], "")
def test_standardize_does_not_blank_default_flat_rate(self):
"""Shared path: default flat rate keeps UOM (blanking is CC-only)."""
constants = Constants()
df = pd.DataFrame(
{
"DEFAULT_IND": ["Y"],
"AARETE_DERIVED_REIMB_METHOD": ["Flat Rate"],
"UNIT_OF_MEASURE": ["Per Unit"],
"REIMB_TERM": [""],
}
)
result = standardize_reimb_method_and_fee_schedule(
df, constants.VALID_UNIT_OF_MEASURE
)
self.assertEqual(result.loc[0, "UNIT_OF_MEASURE"], "Per Unit")
def test_non_default_flat_rate_uom_per_unit(self):
"""DEFAULT_IND='N', flat rate, empty UOM -> UNIT_OF_MEASURE becomes Per Unit."""
constants = Constants()
df = pd.DataFrame(
{
"DEFAULT_IND": ["N"],
"AARETE_DERIVED_REIMB_METHOD": ["Flat Rate"],
"UNIT_OF_MEASURE": [""],
"REIMB_TERM": [""],
}
)
result = standardize_reimb_method_and_fee_schedule(
df, constants.VALID_UNIT_OF_MEASURE
)
self.assertEqual(result.loc[0, "UNIT_OF_MEASURE"], "Per Unit")
def test_default_non_flat_rate_uom_unchanged(self):
"""DEFAULT_IND='Y' and Fee Schedule -> UNIT_OF_MEASURE unchanged."""
constants = Constants()
df = pd.DataFrame(
{
"DEFAULT_IND": ["Y"],
"AARETE_DERIVED_REIMB_METHOD": ["Fee Schedule"],
"UNIT_OF_MEASURE": ["Per Visit"],
"REIMB_TERM": [""],
}
)
result = standardize_reimb_method_and_fee_schedule(
df, constants.VALID_UNIT_OF_MEASURE
)
self.assertEqual(result.loc[0, "UNIT_OF_MEASURE"], "Per Visit")
class TestOutputFileStructure(unittest.TestCase):
"""Test that the new output file structure is generated correctly."""
def setUp(self):
"""Create temporary directory for outputs."""
self.temp_dir = tempfile.mkdtemp()
self.run_timestamp = "run_20250106_10-30_test"
self.batch_id = "test_batch"
def tearDown(self):
"""Clean up temporary directory."""
import shutil
shutil.rmtree(self.temp_dir, ignore_errors=True)
@patch("src.utils.io_utils.config")
def test_write_local_cc_results_full_structure(self, mock_config):
"""Test that cc_results_full creates correct directory structure."""
mock_config.BATCH_ID = "test_batch"
mock_config.CONSOLIDATED_OUTPUT_DIRECTORY = self.temp_dir
mock_config.PERFORM_PARENT_CHILD_MAPPING = False
mock_config.MAX_ROWS_PER_SPLIT = 70000
from src.utils.io_utils import write_local
# Add FILE_NAME column required for splitting logic
test_df = pd.DataFrame(
{"FILE_NAME": ["file1.txt", "file2.txt"], "col1": [1, 2], "col2": [3, 4]}
)
write_local(test_df, "", self.run_timestamp, "cc_results_full")
# Check directory structure was created
expected_dir = os.path.join(
self.temp_dir, self.run_timestamp, "full_outputs", "cc_results"
)
self.assertTrue(os.path.exists(expected_dir))
# Check file was created
expected_file = os.path.join(expected_dir, "test_batch-RESULTS-FULL.csv")
self.assertTrue(os.path.exists(expected_file))
@patch("src.utils.io_utils.config")
def test_write_local_dashboard_results_full_structure(self, mock_config):
"""Test that dashboard_results_full creates correct directory structure."""
mock_config.BATCH_ID = "test_batch"
mock_config.CONSOLIDATED_OUTPUT_DIRECTORY = self.temp_dir
mock_config.PERFORM_PARENT_CHILD_MAPPING = False
mock_config.MAX_ROWS_PER_SPLIT = 70000
from src.utils.io_utils import write_local
# Add FILE_NAME column required for splitting logic
test_df = pd.DataFrame(
{"FILE_NAME": ["file1.txt", "file2.txt"], "col1": [1, 2], "col2": [3, 4]}
)
write_local(test_df, "", self.run_timestamp, "dashboard_results_full")
# Check directory structure was created
expected_dir = os.path.join(
self.temp_dir, self.run_timestamp, "full_outputs", "dashboard_results"
)
self.assertTrue(os.path.exists(expected_dir))
# Check file was created
expected_file = os.path.join(expected_dir, "test_batch-RESULTS-dashboard.csv")
self.assertTrue(os.path.exists(expected_file))
@patch("src.utils.io_utils.config")
def test_write_local_error_in_full_outputs(self, mock_config):
"""Test that error files go to full_outputs/ directory."""
mock_config.BATCH_ID = "test_batch"
mock_config.CONSOLIDATED_OUTPUT_DIRECTORY = self.temp_dir
mock_config.PERFORM_PARENT_CHILD_MAPPING = False
from src.utils.io_utils import write_local
test_df = pd.DataFrame({"error": ["Error 1", "Error 2"]})
write_local(test_df, "", self.run_timestamp, "error")
# Check directory structure was created
expected_dir = os.path.join(self.temp_dir, self.run_timestamp, "full_outputs")
self.assertTrue(os.path.exists(expected_dir))
# Check file was created
expected_file = os.path.join(expected_dir, "test_batch-ERRORS.csv")
self.assertTrue(os.path.exists(expected_file))
@patch("src.utils.io_utils.config")
def test_write_local_qc_qa_cc_full_structure(self, mock_config):
"""Test that qc_qa_cc_full creates correct directory structure."""
mock_config.BATCH_ID = "test_batch"
mock_config.CONSOLIDATED_OUTPUT_DIRECTORY = self.temp_dir
mock_config.PERFORM_PARENT_CHILD_MAPPING = False
mock_config.MAX_ROWS_PER_SPLIT = 70000
from src.utils.io_utils import write_local
# Add FILE_NAME column required for splitting logic
test_df = pd.DataFrame(
{"FILE_NAME": ["file1.txt", "file2.txt"], "col1": [1, 2], "col2": [3, 4]}
)
write_local(test_df, "", self.run_timestamp, "qc_qa_cc_full")
# Check directory structure was created
expected_dir = os.path.join(
self.temp_dir, self.run_timestamp, "automation_qa-qc"
)
self.assertTrue(os.path.exists(expected_dir))
# Check file was created
expected_file = os.path.join(expected_dir, "test_batch-RESULTS-QC-QA-FULL.csv")
self.assertTrue(os.path.exists(expected_file))
@patch("src.utils.io_utils.config")
def test_write_local_qc_qa_stats_structure(self, mock_config):
"""Test that qc_qa_stats creates correct directory structure."""
mock_config.BATCH_ID = "test_batch"
mock_config.CONSOLIDATED_OUTPUT_DIRECTORY = self.temp_dir
mock_config.PERFORM_PARENT_CHILD_MAPPING = False
from src.utils.io_utils import write_local
test_df = pd.DataFrame({"stat": ["value1", "value2"]})
write_local(test_df, "", self.run_timestamp, "qc_qa_stats")
# Check directory structure was created
expected_dir = os.path.join(
self.temp_dir, self.run_timestamp, "automation_qa-qc"
)
self.assertTrue(os.path.exists(expected_dir))
# Check file was created
expected_file = os.path.join(expected_dir, "test_batch-QC-QA-STATS.csv")
self.assertTrue(os.path.exists(expected_file))
@patch("src.utils.io_utils.config")
def test_write_local_parent_child_structure(self, mock_config):
"""Test that parent_child creates correct directory structure and file path."""
mock_config.BATCH_ID = "test_batch"
mock_config.CONSOLIDATED_OUTPUT_DIRECTORY = self.temp_dir
mock_config.PERFORM_PARENT_CHILD_MAPPING = False
from src.utils.io_utils import write_local
test_df = pd.DataFrame({"col1": [1, 2]})
# Mock to_csv on the DataFrame instance to avoid actual file I/O
test_df.to_csv = MagicMock()
write_local(test_df, "", self.run_timestamp, "parent_child")
# Verify to_csv was called on the DataFrame
test_df.to_csv.assert_called_once()
# Verify the call had index=False and quoting=1
call_kwargs = test_df.to_csv.call_args[1]
self.assertEqual(call_kwargs.get("index"), False)
self.assertEqual(call_kwargs.get("quoting"), 1)
# Verify the file path contains the expected components
call_args = test_df.to_csv.call_args[0]
file_path = call_args[0] if call_args else None
self.assertIsNotNone(file_path, "to_csv should be called with a file path")
self.assertIn(
"parent-child",
file_path,
f"File path should contain 'parent-child', got: {file_path}",
)
self.assertIn(
"test_batch-PC.csv",
file_path,
f"File path should contain 'test_batch-PC.csv', got: {file_path}",
)
# Verify the directory structure would be correct
expected_dir = os.path.join(self.temp_dir, self.run_timestamp, "parent-child")
self.assertIn(
expected_dir,
file_path,
f"File path should contain expected directory {expected_dir}, got: {file_path}",
)
@patch("src.utils.io_utils.config")
@patch("src.utils.io_utils.logging")
def test_write_s3_cc_results_full_path(self, mock_logging, mock_config):
"""Test that write_s3 uses correct S3 path for cc_results_full."""
mock_config.BATCH_ID = "test_batch"
mock_config.S3_OUTPUT_BUCKET = "test-bucket"
mock_config.S3_CLIENT = MagicMock()
mock_config.S3_CLIENT.put_object = MagicMock()
mock_config.PERFORM_PARENT_CHILD_MAPPING = False
mock_config.MAX_ROWS_PER_SPLIT = 70000
from src.utils.io_utils import write_s3
# Add FILE_NAME column required for splitting logic
test_df = pd.DataFrame(
{"FILE_NAME": ["file1.txt", "file2.txt"], "col1": [1, 2]}
)
write_s3(test_df, "", self.run_timestamp, "cc_results_full")
# Verify S3 put_object was called
mock_config.S3_CLIENT.put_object.assert_called_once()
# Verify the S3 key/path
call_args = mock_config.S3_CLIENT.put_object.call_args
s3_key = call_args[1]["Key"]
expected_key = f"test_batch/{self.run_timestamp}/full_outputs/cc_results/test_batch-RESULTS-FULL.csv"
self.assertEqual(s3_key, expected_key)
@patch("src.utils.io_utils.config")
@patch("src.utils.io_utils.logging")
def test_write_s3_dashboard_results_full_path(self, mock_logging, mock_config):
"""Test that write_s3 uses correct S3 path for dashboard_results_full."""
mock_config.BATCH_ID = "test_batch"
mock_config.S3_OUTPUT_BUCKET = "test-bucket"
mock_config.S3_CLIENT = MagicMock()
mock_config.S3_CLIENT.put_object = MagicMock()
mock_config.PERFORM_PARENT_CHILD_MAPPING = False
mock_config.MAX_ROWS_PER_SPLIT = 70000
from src.utils.io_utils import write_s3
# Add FILE_NAME column required for splitting logic
test_df = pd.DataFrame(
{"FILE_NAME": ["file1.txt", "file2.txt"], "col1": [1, 2]}
)
write_s3(test_df, "", self.run_timestamp, "dashboard_results_full")
# Verify S3 put_object was called
mock_config.S3_CLIENT.put_object.assert_called_once()
# Verify the S3 key/path
call_args = mock_config.S3_CLIENT.put_object.call_args
s3_key = call_args[1]["Key"]
expected_key = f"test_batch/{self.run_timestamp}/full_outputs/dashboard_results/test_batch-RESULTS-dashboard.csv"
self.assertEqual(s3_key, expected_key)
@patch("src.utils.io_utils.config")
@patch("src.utils.io_utils.logging")
def test_write_s3_qc_qa_paths(self, mock_logging, mock_config):
"""Test that write_s3 uses correct S3 paths for QC/QA outputs."""
mock_config.BATCH_ID = "test_batch"
mock_config.S3_OUTPUT_BUCKET = "test-bucket"
mock_config.S3_CLIENT = MagicMock()
mock_config.S3_CLIENT.put_object = MagicMock()
mock_config.PERFORM_PARENT_CHILD_MAPPING = False
mock_config.MAX_ROWS_PER_SPLIT = 70000
from src.utils.io_utils import write_s3
# Add FILE_NAME column required for splitting logic
test_df = pd.DataFrame(
{"FILE_NAME": ["file1.txt", "file2.txt"], "col1": [1, 2]}
)
# Test qc_qa_cc_full
write_s3(test_df, "", self.run_timestamp, "qc_qa_cc_full")
call_args = mock_config.S3_CLIENT.put_object.call_args
s3_key = call_args[1]["Key"]
expected_key = f"test_batch/{self.run_timestamp}/automation_qa-qc/test_batch-RESULTS-QC-QA-FULL.csv"
self.assertEqual(s3_key, expected_key)
# Test qc_qa_stats
mock_config.S3_CLIENT.put_object.reset_mock()
test_df_stats = pd.DataFrame(
{"col1": [1, 2]}
) # qc_qa_stats doesn't need FILE_NAME
write_s3(test_df_stats, "", self.run_timestamp, "qc_qa_stats")
call_args = mock_config.S3_CLIENT.put_object.call_args
s3_key = call_args[1]["Key"]
expected_key = f"test_batch/{self.run_timestamp}/automation_qa-qc/test_batch-QC-QA-STATS.csv"
self.assertEqual(s3_key, expected_key)
if __name__ == "__main__":
unittest.main()
Reference in New Issue View Git Blame Copy Permalink