doczyai-pipelines/fieldExtraction/tests/test_tin_npi_funcs.py

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"]