doczyai-pipelines/fieldExtraction/tests/test_tin_npi_funcs.py

import pytest
from src.investment.tin_npi_funcs import clean_provider_info, get_all_matches, chunk_on_matches

def test_clean_provider_info_valid_data():
    provider_info = [
        {"TIN": "12-3456789", "NPI": "1234567890", "NAME": "John Doe", "IS_GROUP": True},
        {"TIN": "98.7654321", "NPI": "0987654321", "NAME": "  Jane   Smith  ", "IS_GROUP": False}
    ]
    expected_output = [
        {"TIN": "123456789", "NPI": "1234567890", "NAME": "John Doe", "IS_GROUP": True},
        {"TIN": "987654321", "NPI": "0987654321", "NAME": "Jane Smith", "IS_GROUP": False}
    ]
    assert clean_provider_info(provider_info) == expected_output

def test_clean_provider_info_invalid_tin_npi():
    provider_info = [
        {"TIN": "12345", "NPI": "12345", "NAME": "John Doe", "IS_GROUP": True},
        {"TIN": None, "NPI": None, "NAME": "Jane Smith", "IS_GROUP": False}
    ]
    expected_output = [
        {"TIN": "UNKNOWN", "NPI": "UNKNOWN", "NAME": "John Doe", "IS_GROUP": True},
        {"TIN": "UNKNOWN", "NPI": "UNKNOWN", "NAME": "Jane Smith", "IS_GROUP": False}
    ]
    assert clean_provider_info(provider_info) == expected_output

def test_clean_provider_info_missing_fields():
    provider_info = [
        {"TIN": "12-3456789", "NAME": "John Doe"},
        {"NPI": "1234567890", "IS_GROUP": True}
    ]
    expected_output = [
        {"TIN": "123456789", "NPI": "UNKNOWN", "NAME": "John Doe", "IS_GROUP": False},
        {"TIN": "UNKNOWN", "NPI": "1234567890", "NAME": "UNKNOWN", "IS_GROUP": True}
    ]
    assert clean_provider_info(provider_info) == expected_output

def test_clean_provider_info_empty_input():
    provider_info = []
    expected_output = []
    assert clean_provider_info(provider_info) == expected_output

def test_clean_provider_info_whitespace_name():
    provider_info = [
        {"TIN": "12-3456789", "NPI": "1234567890", "NAME": "   ", "IS_GROUP": True}
    ]
    expected_output = [
        {"TIN": "123456789", "NPI": "1234567890", "NAME": "UNKNOWN", "IS_GROUP": True}
    ]
    assert clean_provider_info(provider_info) == expected_output

def test_get_all_matches_single_match():
    text = "The TIN is 123456789."
    pattern = r"\b\d{9}\b"
    filename = "test_file.txt"
    expected_output = ["123456789"]
    assert get_all_matches(text, pattern, filename) == expected_output

def test_get_all_matches_multiple_matches():
    text = "TINs: 123456789, 987654321, and 123456789."
    pattern = r"\b\d{9}\b"
    filename = "test_file.txt"
    expected_output = ["123456789", "987654321"]
    assert sorted(get_all_matches(text, pattern, filename)) == sorted(expected_output)

def test_get_all_matches_no_matches():
    text = "No valid TINs here."
    pattern = r"\b\d{9}\b"
    filename = "test_file.txt"
    expected_output = []
    assert get_all_matches(text, pattern, filename) == expected_output

def test_get_all_matches_special_characters():
    text = "TINs: 123-45-6789 and 987.65.4321."
    pattern = r"\b\d{3}-\d{2}-\d{4}\b|\b\d{3}\.\d{2}\.\d{4}\b"
    filename = "test_file.txt"
    expected_output = ["123-45-6789", "987.65.4321"]
    assert sorted(get_all_matches(text, pattern, filename)) == sorted(expected_output)

def test_get_all_matches_empty_text():
    text = ""
    pattern = r"\b\d{9}\b"
    filename = "test_file.txt"
    expected_output = []
    assert get_all_matches(text, pattern, filename) == expected_output

def test_get_all_matches_empty_pattern():
    text = "The TIN is 123456789."
    pattern = ""
    filename = "test_file.txt"
    expected_output = []
    assert get_all_matches(text, pattern, filename) == expected_output

def test_chunk_on_matches_single_match():
    matches = ["match1"]
    text_dict = {
        "page1": "This is a text with match1.",
        "page2": "This is another text without matches."
    }
    expected_output = "This is a text with match1."
    assert chunk_on_matches(matches, text_dict) == expected_output

def test_chunk_on_matches_multiple_matches():
    matches = ["match1", "match2"]
    text_dict = {
        "page1": "This is a text with match1.",
        "page2": "This is another text with match2.",
        "page3": "No matches here."
    }
    expected_output = "This is a text with match1.\nThis is another text with match2."
    assert chunk_on_matches(matches, text_dict) == expected_output

def test_chunk_on_matches_no_matches():
    matches = ["match1"]
    text_dict = {
        "page1": "No relevant text here.",
        "page2": "Still no matches."
    }
    expected_output = ""
    assert chunk_on_matches(matches, text_dict) == expected_output

def test_chunk_on_matches_empty_matches():
    matches = []
    text_dict = {
        "page1": "This is some text.",
        "page2": "This is more text."
    }
    expected_output = ""
    assert chunk_on_matches(matches, text_dict) == expected_output

def test_chunk_on_matches_empty_text_dict():
    matches = ["match1"]
    text_dict = {}
    expected_output = ""
    assert chunk_on_matches(matches, text_dict) == expected_output

def test_chunk_on_matches_partial_match():
    matches = ["match"]
    text_dict = {
        "page1": "This is a text with match1.",
        "page2": "This is another text with match2.",
        "page3": "But this page you won't find one."
    }
    expected_output = "This is a text with match1.\nThis is another text with match2."
    assert chunk_on_matches(matches, text_dict) == expected_output

def test_chunk_on_matches_whitespace_handling():
    matches = ["match1"]
    text_dict = {
        "page1": "   This is a text with match1.   ",
        "page2": "This is another text without matches."
    }
    expected_output = "   This is a text with match1.   "
    assert chunk_on_matches(matches, text_dict) == expected_output