Katon Minhas
45db295d17
derive IS_GROUP * derive IS_GROUP * Update unit tests Approved-by: Alex Galarce
149 lines
4.9 KiB
Python
149 lines
4.9 KiB
Python
import pytest
|
|
from unittest.mock import patch, MagicMock
|
|
import json
|
|
|
|
import src.investment.tin_npi_funcs as tin_npi_funcs
|
|
from src.prompts.investment_prompts 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, "test.pdf")
|
|
assert sorted(result) == sorted(['12-3456789', '98-7654321'])
|
|
|
|
# Test empty pattern
|
|
assert tin_npi_funcs.get_all_matches(text, "", "test.pdf") == []
|
|
|
|
# Test no matches
|
|
assert tin_npi_funcs.get_all_matches("No TINs here", pattern, "test.pdf") == []
|
|
|
|
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"
|
|
|
|
def test_merge_provider_info(self):
|
|
one_to_one_results = {}
|
|
provider_info = [
|
|
{
|
|
"TIN": "123456789",
|
|
"NPI": "1234567890",
|
|
"NAME": "Group Provider",
|
|
"IS_GROUP": "Y"
|
|
},
|
|
{
|
|
"TIN": "987654321",
|
|
"NPI": "9876543210",
|
|
"NAME": "Other Provider",
|
|
"IS_GROUP": "N"
|
|
}
|
|
]
|
|
|
|
result = tin_npi_funcs.merge_provider_info(one_to_one_results, provider_info)
|
|
|
|
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"]
|
|
|
|
|
|
def test_deduplicate_providers(self):
|
|
providers = [
|
|
{
|
|
"TIN": "123456789",
|
|
"NPI": "1234567890",
|
|
"NAME": "Provider A"
|
|
},
|
|
{
|
|
"TIN": "123456789",
|
|
"NPI": "1234567890",
|
|
"NAME": "Provider A"
|
|
},
|
|
{
|
|
"TIN": "UNKNOWN",
|
|
"NPI": "UNKNOWN",
|
|
"NAME": "UNKNOWN"
|
|
}
|
|
]
|
|
|
|
result = tin_npi_funcs.deduplicate_providers(providers)
|
|
assert len(result) == 1
|
|
assert result[0]["TIN"] == "123456789"
|
|
|
|
@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",
|
|
"ON_SIGNATURE_PAGE": "N"
|
|
}])
|
|
|
|
result = tin_npi_funcs.get_provider_info(sample_text_dict, "1", "test.pdf")
|
|
assert len(result) == 1
|
|
assert result[0]["IS_GROUP"] == "Y" # Page 1 should be marked as intro
|
|
|
|
@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"
|
|
}])
|
|
|
|
one_to_one_fields = FieldSet(relationship="one_to_one")
|
|
contract_text = "Contract with TIN 12-3456789"
|
|
|
|
results, fields = tin_npi_funcs.run_provider_info_fields(
|
|
contract_text,
|
|
one_to_one_fields,
|
|
sample_text_dict,
|
|
"test.pdf"
|
|
)
|
|
|
|
assert "PROV_INFO_JSON" in results
|
|
assert "PROV_INFO_JSON_FORMATTED" in results
|
|
assert isinstance(results["PROV_INFO_JSON"], str) |