Files
doczyai-pipelines/fieldExtraction/tests/test_tin_npi_funcs.py
T
VenkataKrishna Reddy Avula d97a31aa5b Merged in optimization/prov_info (pull request #829)
optimization_prov_info to main

* removed on signature page field which is unnecessary


Approved-by: Katon Minhas
2026-01-07 15:19:35 +00:00

341 lines
12 KiB
Python

import json
import re
from unittest.mock import MagicMock, patch
import constants.regex_patterns as regex_patterns
import pytest
import src.investment.tin_npi_funcs as tin_npi_funcs
from src.prompts.fieldset import FieldSet
class TestTinNpiFuncs:
@pytest.fixture
def sample_text_dict(self):
return {
"1": "First page with TIN 12-3456789",
"2": "Second page with NPI 1234567890",
"2.1": "Continuation with same TIN 12-3456789",
"3": "Page with both TIN 98-7654321 and NPI 9876543210",
}
def test_get_all_matches(self):
text = "TIN: 12-3456789 and another TIN: 98-7654321"
pattern = r"\d{2}[-\s]?\d{7}"
result = tin_npi_funcs.get_all_matches(text, pattern)
assert sorted(result) == sorted(["12-3456789", "98-7654321"])
# Test empty pattern
assert tin_npi_funcs.get_all_matches(text, "") == []
# Test no matches
assert tin_npi_funcs.get_all_matches("No TINs here", pattern) == []
def test_chunk_on_matches(self):
text_dict = {
"1": "Page with match1",
"2": "Page without matches",
"3": "Page with match2",
}
matches = ["match1", "match2"]
result = tin_npi_funcs.chunk_on_matches(matches, text_dict)
assert "match1" in result
assert "match2" in result
assert "without matches" not in result
def test_clean_provider_info(self):
providers = [
{
"TIN": "12-345.6789",
"NPI": "12.34567890",
"NAME": "Test Provider",
"IS_GROUP": "Y",
},
{
"TIN": "",
"NPI": "invalid",
"NAME": "",
"IS_GROUP": "N",
},
]
result = tin_npi_funcs.clean_provider_info(providers)
assert result[0]["TIN"] == "123456789"
assert result[0]["NPI"] == "1234567890"
assert result[0]["NAME"] == "Test Provider"
assert result[1]["TIN"] == "UNKNOWN"
assert result[1]["NPI"] == "UNKNOWN"
assert result[1]["NAME"] == "UNKNOWN"
@patch("src.investment.prompt_calls.provider_name_match_check")
def test_merge_provider_info_with_hybrid_smart_chunking(self, mock_name_match_check):
"""Test merge_provider_info_with_hybrid_smart_chunking with name matching logic"""
# Mock the name matching to return True for "Group Provider"
def name_match_side_effect(provider_name, name, filename):
return name == "Group Provider"
mock_name_match_check.side_effect = name_match_side_effect
one_to_one_results = {
"PROVIDER_NAME": "Group Provider",
"PROV_INFO_JSON": json.dumps([
{
"TIN": "123456789",
"NPI": "1234567890",
"NAME": "Group Provider",
},
{
"TIN": "987654321",
"NPI": "9876543210",
"NAME": "Other Provider",
},
])
}
result = tin_npi_funcs.merge_provider_info_with_hybrid_smart_chunking(one_to_one_results, "test.pdf")
# Check that IS_GROUP flags were added correctly
prov_info = json.loads(result["PROV_INFO_JSON"])
assert prov_info[0]["IS_GROUP"] == "Y"
assert prov_info[1]["IS_GROUP"] == "N"
# Check group and other fields
assert result["PROV_GROUP_TIN"] == "123456789"
assert result["PROV_OTHER_TIN"] == "987654321"
assert "Group Provider" in result["PROV_GROUP_NAME_FULL"]
assert "Other Provider" in result["PROV_OTHER_NAME_FULL"]
@patch("src.utils.llm_utils.invoke_claude")
def test_get_provider_info(self, mock_invoke_claude, sample_text_dict):
mock_invoke_claude.return_value = json.dumps(
[
{
"TIN": "123456789",
"NPI": "1234567890",
"NAME": "Test Provider"
}
]
)
result = tin_npi_funcs.get_provider_info(sample_text_dict, "1", "test.pdf", payer_name="Test Payer")
assert len(result) == 1
@patch("src.utils.llm_utils.invoke_claude")
def test_run_provider_info_fields(self, mock_invoke_claude, sample_text_dict):
mock_invoke_claude.return_value = json.dumps(
[
{
"TIN": "123456789",
"NPI": "1234567890",
"NAME": "Test Provider",
}
]
)
contract_text = "Contract with TIN 12-3456789"
results = tin_npi_funcs.run_provider_info_fields(
contract_text, sample_text_dict, "test.pdf", "Test Payer"
)
assert "PROV_INFO_JSON" in results
assert "PROV_INFO_JSON_FORMATTED" in results
assert isinstance(results["PROV_INFO_JSON"], str)
def test_get_all_matches_with_ocr_exact_matches(self):
"""Test that exact matches are found and marked correctly"""
text = "TIN: 123456789 and NPI: 1234567890"
tin_results = tin_npi_funcs.get_all_matches_with_ocr(
text, r"\b\d{9}\b", r"\b[O0lI1S5G6B8gbo\d]{9}\b", "TIN"
)
assert len(tin_results) == 1
assert tin_results[0] == ("123456789", "EXACT")
def test_get_all_matches_with_ocr_corrected_matches(self):
"""Test that OCR-corrected matches work properly"""
text = "TIN: 12345678O and Provider ID: l234567890"
tin_results = tin_npi_funcs.get_all_matches_with_ocr(
text, r"\b\d{9}\b", r"\b[O0lI1S5G6B8gbo\d]{9}\b", "TIN"
)
npi_results = tin_npi_funcs.get_all_matches_with_ocr(
text, r"\b\d{10}\b", r"\b[O0lI1S5G6B8gbo\d]{10}\b", "NPI"
)
assert len(tin_results) == 1
assert tin_results[0] == ("123456780", "OCR_CORRECTED")
assert len(npi_results) == 1
assert npi_results[0] == ("1234567890", "OCR_CORRECTED")
def test_get_all_matches_with_ocr_mixed_matches(self):
"""Test that both exact and OCR matches are found in same text"""
text = "Clean TIN: 123456789 and corrupted TIN: 98765432O"
results = tin_npi_funcs.get_all_matches_with_ocr(
text, r"\b\d{9}\b", r"\b[O0lI1S5G6B8gbo\d]{9}\b", "TIN"
)
assert len(results) == 2
# Sort results to ensure consistent order
results.sort(key=lambda x: x[0])
assert results[0] == ("123456789", "EXACT")
assert results[1] == ("987654320", "OCR_CORRECTED")
def test_get_all_matches_with_ocr_no_duplicates(self):
"""Test that OCR correction doesn't create duplicates of exact matches"""
text = "TIN: 123456789 and same TIN with OCR error: 12345G789"
results = tin_npi_funcs.get_all_matches_with_ocr(
text, r"\b\d{9}\b", r"\b[O0lI1S5G6B8gbo\d]{9}\b", "TIN"
)
# Should only get one result (the exact match wins)
assert len(results) == 1
assert results[0] == ("123456789", "EXACT")
def test_attempt_ocr_correction_tin_success(self):
"""Test successful TIN OCR correction"""
result, is_valid = tin_npi_funcs.attempt_ocr_correction("12345678O", "TIN")
assert is_valid is True
assert result == "123456780"
result, is_valid = tin_npi_funcs.attempt_ocr_correction("l23456789", "TIN")
assert is_valid is True
assert result == "123456789"
def test_attempt_ocr_correction_npi_success(self):
"""Test successful NPI OCR correction"""
result, is_valid = tin_npi_funcs.attempt_ocr_correction("123456789O", "NPI")
assert is_valid is True
assert result == "1234567890"
result, is_valid = tin_npi_funcs.attempt_ocr_correction("l234567890", "NPI")
assert is_valid is True
assert result == "1234567890"
def test_attempt_ocr_correction_failure_cases(self):
"""Test cases where OCR correction should fail"""
# Wrong length after correction
result, is_valid = tin_npi_funcs.attempt_ocr_correction("12345678", "TIN")
assert is_valid is False
assert result == "12345678"
# Invalid characters that aren't in OCR substitution map
result, is_valid = tin_npi_funcs.attempt_ocr_correction("123ABC789", "TIN")
assert is_valid is False
assert result == "123ABC789"
# Unsupported identifier type
result, is_valid = tin_npi_funcs.attempt_ocr_correction("123456789", "UNKNOWN")
assert is_valid is False
assert result == "123456789"
def test_attempt_ocr_correction_formatted_tins(self):
"""Test OCR correction on formatted TINs"""
result, is_valid = tin_npi_funcs.attempt_ocr_correction("12-345678O", "TIN")
assert is_valid is True
assert result == "123456780"
result, is_valid = tin_npi_funcs.attempt_ocr_correction("123-4S-6789", "TIN")
assert is_valid is True
assert result == "123456789"
def test_attempt_ocr_correction_multiple_substitutions(self):
"""Test OCR correction with multiple character substitutions"""
result, is_valid = tin_npi_funcs.attempt_ocr_correction("l2345678O", "TIN")
assert is_valid is True
assert result == "123456780"
result, is_valid = tin_npi_funcs.attempt_ocr_correction("lS345678O", "TIN")
assert is_valid is True
assert result == "153456780"
@patch("src.utils.llm_utils.invoke_claude")
def test_run_provider_info_fields_with_ocr_quality_tracking(
self, mock_invoke_claude, sample_text_dict
):
"""Test that quality tracking fields are added to results"""
mock_invoke_claude.return_value = json.dumps(
[
{
"TIN": "123456789",
"NPI": "1234567890",
"NAME": "Test Provider",
}
]
)
# Include both exact and OCR-correctable TINs
contract_text = "Contract with TIN 12-3456789 and corrupted TIN 98765432O"
results = tin_npi_funcs.run_provider_info_fields(
contract_text, sample_text_dict, "test.pdf", "Test Pyer"
)
# Check that quality tracking fields are present
assert "TIN_EXTRACTION_QUALITY" in results
assert "NPI_EXTRACTION_QUALITY" in results
# Check format of quality strings
assert "exact" in results["TIN_EXTRACTION_QUALITY"]
assert "OCR-corrected" in results["TIN_EXTRACTION_QUALITY"]
def test_ocr_regex_patterns(self):
"""Test that OCR regex patterns match expected cases"""
# Test TIN OCR pattern
tin_ocr_cases = [
"12345678O", # O at end
"O23456789", # O at start
"1234O6789", # O in middle
"12-345678O", # Formatted with O
"123-4S-6789", # Formatted with S
]
for case in tin_ocr_cases:
matches = re.findall(regex_patterns.TIN_OCR_PATTERN, case)
assert len(matches) == 1, f"TIN OCR pattern should match '{case}'"
# Test NPI OCR pattern
npi_ocr_cases = [
"123456789O", # O at end
"l234567890", # l at start
"12345O7890", # O in middle
]
for case in npi_ocr_cases:
matches = re.findall(regex_patterns.NPI_OCR_PATTERN, case)
assert len(matches) == 1, f"NPI OCR pattern should match '{case}'"
@patch("src.utils.llm_utils.invoke_claude")
def test_run_provider_info_fields_no_exact_matches_but_ocr_matches(
self, mock_invoke_claude, sample_text_dict
):
"""Test scenario where no exact matches found but OCR matches are"""
mock_invoke_claude.return_value = json.dumps(
[
{
"TIN": "123456780", # Corrected from 12345678O
"NPI": "1234567890", # Corrected from 123456789O
"NAME": "Test Provider",
}
]
)
# Only OCR-corrupted identifiers, no clean ones
contract_text = "Contract with corrupted TIN 12345678O and NPI 123456789O"
results = tin_npi_funcs.run_provider_info_fields(
contract_text, sample_text_dict, "test.pdf", "Test Payer"
)
# Should find OCR-corrected values
assert "0 exact, 1 OCR-corrected" in results["TIN_EXTRACTION_QUALITY"]
assert "0 exact, 1 OCR-corrected" in results["NPI_EXTRACTION_QUALITY"]