doczyai-pipelines/src/tests/test_tin_npi_funcs.py

import json
import re
from unittest.mock import MagicMock, patch

import src.constants.regex_patterns as regex_patterns
import pytest
import src.pipelines.shared.extraction.tin_npi_funcs as tin_npi_funcs
from src.prompts.fieldset import FieldSet
from src import config


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) == []

    @patch("src.utils.llm_utils.invoke_claude")
    def test_get_provider_info(self, mock_invoke_claude, sample_text_dict):
        """Test prompt_provider_info function (renamed from get_provider_info)."""
        from src.pipelines.saas.prompts import prompt_calls

        mock_invoke_claude.return_value = json.dumps(
            [{"TIN": "123456789", "NPI": "1234567890", "NAME": "Test Provider"}]
        )

        provider_fields = FieldSet(
            file_path=config.FIELD_JSON_PATH, field_type="provider_info"
        )
        result = prompt_calls.prompt_provider_info(
            sample_text_dict,
            "1",
            "test.pdf",
            payer_name="Test Payer",
            provider_fields=provider_fields,
        )
        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",
                }
            ]
        )

        results = tin_npi_funcs.run_provider_info_fields(
            sample_text_dict, "test.pdf", "Test Payer"
        )

        assert "PROV_INFO_JSON" in results
        # PROV_INFO_JSON is now a list, not a JSON string
        assert isinstance(results["PROV_INFO_JSON"], list)

    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(
            sample_text_dict, "test.pdf", "Test Pyer"
        )

    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(
            sample_text_dict, "test.pdf", "Test Payer"
        )

    def test_deduplicate_providers_case_insensitive_names(self):
        """Test that provider deduplication is case-insensitive for names."""
        provider_info_json = [
            {"TIN": "123456789", "NPI": "1234567890", "NAME": "Test Provider"},
            {"TIN": "123456789", "NPI": "1234567890", "NAME": "test provider"},
            {"TIN": "123456789", "NPI": "1234567890", "NAME": "TEST PROVIDER"},
        ]

        result = tin_npi_funcs.deduplicate_providers(provider_info_json)

        # Should only have one provider since all fields match (case-insensitive for name)
        assert len(result) == 1
        assert result[0]["NAME"] == "Test Provider"

    def test_deduplicate_providers_different_cases(self):
        """Test deduplication with different TINs and case variations in names."""
        provider_info_json = [
            {"TIN": "123456789", "NPI": "N/A", "NAME": "Provider One"},
            {"TIN": "123456789", "NPI": "N/A", "NAME": "provider one"},
            {"TIN": "987654321", "NPI": "N/A", "NAME": "provider one"},
            {"TIN": "111222333", "NPI": "N/A", "NAME": "Provider Two"},
        ]

        result = tin_npi_funcs.deduplicate_providers(provider_info_json)

        # Should have 3 providers:
        # - Row 1 and 2 are duplicates (same TIN, same name case-insensitively) -> keep Row 1
        # - Row 3 is unique (different TIN)
        # - Row 4 is unique
        assert len(result) == 3

    def test_add_filename_tin_merge_single_provider_single_tin(self):
        """Test merge logic when there's 1 provider with N/A TIN and 1 filename TIN."""
        one_to_one_results = {
            "PROV_INFO_JSON": [
                {"TIN": "N/A", "NPI": "1234567890", "NAME": "Test Provider"}
            ]
        }
        filename = "contract-123456789-test.txt"

        result = tin_npi_funcs.add_filename_tin(one_to_one_results, filename)

        # Should merge the filename TIN into the existing provider
        assert len(result["PROV_INFO_JSON"]) == 1
        assert result["PROV_INFO_JSON"][0]["TIN"] == "123456789"
        assert result["PROV_INFO_JSON"][0]["FROM_FILENAME"] == "Y"
        assert result["PROV_INFO_JSON"][0]["NPI"] == "1234567890"

    def test_add_filename_tin_no_merge_multiple_providers(self):
        """Test that merge doesn't happen with multiple providers."""
        one_to_one_results = {
            "PROV_INFO_JSON": [
                {"TIN": "N/A", "NPI": "1234567890", "NAME": "Provider 1"},
                {"TIN": "987654321", "NPI": "N/A", "NAME": "Provider 2"},
            ]
        }
        filename = "contract-123456789-test.txt"

        result = tin_npi_funcs.add_filename_tin(one_to_one_results, filename)

        # Should add as separate provider, not merge
        assert len(result["PROV_INFO_JSON"]) == 3
        assert result["PROV_INFO_JSON"][2]["TIN"] == "123456789"
        assert result["PROV_INFO_JSON"][2]["FROM_FILENAME"] == "Y"

    def test_add_filename_tin_no_merge_multiple_filename_tins(self):
        """Test that merge doesn't happen with multiple filename TINs."""
        one_to_one_results = {
            "PROV_INFO_JSON": [
                {"TIN": "N/A", "NPI": "1234567890", "NAME": "Test Provider"}
            ]
        }
        filename = "contract-123456789-987654321.txt"

        result = tin_npi_funcs.add_filename_tin(one_to_one_results, filename)

        # Should add both TINs as separate providers, not merge
        assert len(result["PROV_INFO_JSON"]) == 3
        # Sort both lists for comparison since order may not be deterministic
        assert sorted(result["FILENAME_TIN"]) == sorted(["123456789", "987654321"])

    def test_add_group_fallback_when_empty_and_name_str(self):
        """Empty PROV_INFO_JSON with a string PROVIDER_NAME synthesizes one entry."""
        one_to_one_results = {
            "PROV_INFO_JSON": [],
            "PROVIDER_NAME": "Oregon Health & Science University",
        }

        result = tin_npi_funcs.add_group_and_other(one_to_one_results, "test.txt")

        assert len(result["PROV_INFO_JSON"]) == 1
        entry = result["PROV_INFO_JSON"][0]
        assert entry == {
            "TIN": "N/A",
            "NPI": "N/A",
            "NAME": "Oregon Health & Science University",
            "FROM_FILENAME": "N",
            "IS_GROUP": "Y",
        }
        assert result["PROV_GROUP_NAME_FULL"] == ["Oregon Health & Science University"]
        assert result["PROV_GROUP_TIN"] == []
        assert result["PROV_GROUP_NPI"] == []
        assert result["PROV_OTHER_TIN"] == []
        assert result["PROV_OTHER_NPI"] == []
        assert result["PROV_OTHER_NAME_FULL"] == []

    def test_add_group_fallback_when_empty_and_name_list(self):
        """List PROVIDER_NAME produces one synthesized entry per non-empty name."""
        one_to_one_results = {
            "PROV_INFO_JSON": [],
            "PROVIDER_NAME": ["OHSU", "OHSU Home Infusion Pharmacy"],
        }

        result = tin_npi_funcs.add_group_and_other(one_to_one_results, "test.txt")

        assert len(result["PROV_INFO_JSON"]) == 2
        names = [e["NAME"] for e in result["PROV_INFO_JSON"]]
        assert names == ["OHSU", "OHSU Home Infusion Pharmacy"]
        for entry in result["PROV_INFO_JSON"]:
            assert entry["TIN"] == "N/A"
            assert entry["NPI"] == "N/A"
            assert entry["IS_GROUP"] == "Y"
            assert entry["FROM_FILENAME"] == "N"
        assert result["PROV_GROUP_NAME_FULL"] == [
            "OHSU",
            "OHSU Home Infusion Pharmacy",
        ]

    @patch(
        "src.pipelines.saas.prompts.prompt_calls.provider_name_match_check",
        return_value=True,
    )
    def test_add_group_no_fallback_when_prov_info_already_populated(self, _mock_check):
        """When PROV_INFO_JSON already has entries, the fallback must not fire."""
        one_to_one_results = {
            "PROV_INFO_JSON": [
                {
                    "TIN": "262998718",
                    "NPI": "1073054714",
                    "NAME": "OHSU",
                    "FROM_FILENAME": "N",
                }
            ],
            "PROVIDER_NAME": "OHSU",
        }

        result = tin_npi_funcs.add_group_and_other(one_to_one_results, "test.txt")

        # No synthesized entry: the single original entry remains, and the
        # normal loop processed it (IS_GROUP now set).
        assert len(result["PROV_INFO_JSON"]) == 1
        assert result["PROV_INFO_JSON"][0]["TIN"] == "262998718"
        assert result["PROV_INFO_JSON"][0]["NPI"] == "1073054714"
        assert result["PROV_INFO_JSON"][0]["IS_GROUP"] == "Y"

    @pytest.mark.parametrize("empty_name", ["", None, "N/A", []])
    def test_add_group_no_fallback_when_name_empty(self, empty_name):
        """Empty/None/N/A/[] PROVIDER_NAME leaves PROV_INFO_JSON empty."""
        one_to_one_results = {
            "PROV_INFO_JSON": [],
            "PROVIDER_NAME": empty_name,
        }

        result = tin_npi_funcs.add_group_and_other(one_to_one_results, "test.txt")

        assert result["PROV_INFO_JSON"] == []
        assert result["PROV_GROUP_NAME_FULL"] == []
        assert result["PROV_GROUP_TIN"] == []
        assert result["PROV_GROUP_NPI"] == []

    def test_add_group_no_fallback_when_name_list_all_empty(self):
        """List of only empty/N/A names leaves PROV_INFO_JSON empty."""
        one_to_one_results = {
            "PROV_INFO_JSON": [],
            "PROVIDER_NAME": ["", "N/A", None],
        }

        result = tin_npi_funcs.add_group_and_other(one_to_one_results, "test.txt")

        assert result["PROV_INFO_JSON"] == []
        assert result["PROV_GROUP_NAME_FULL"] == []