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", "ON_SIGNATURE_PAGE": "N", } ] ) 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", "ON_SIGNATURE_PAGE": "N", } ] ) # 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", "ON_SIGNATURE_PAGE": "N", } ] ) # 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"]