doczyai-pipelines/fieldExtraction/tests/test_string_utils.py

import numpy as np
import pandas as pd
import pytest

import src.utils as utils
import src.utils.llm_utils as llm_utils
from constants.delimiters import Delimiter
from src.utils.string_utils import (
    contains_reimbursement,
    count_reimbursements_in_exhibit,
    extract_signature_page,
    extract_text_from_delimiters,
    is_empty,
    json_parsing_search,
    page_key_sort,
)


class TestStringUtils:
    @pytest.mark.parametrize(
        "raw_text, delimiter, match_index, expected",
        [
            (
                "Here is the answer |this is a pipe answer| and |more text.|",
                Delimiter.PIPE,
                0,
                "this is a pipe answer",
            ),
            (
                "Here is the answer |this is a pipe answer| and |more text.|",
                Delimiter.PIPE,
                -1,
                "more text.",
            ),
            (
                "|Here| is the answer |this is a pipe answer| and |more text.|",
                Delimiter.PIPE,
                1,
                "this is a pipe answer",
            ),
            (
                "Here is the answer `this is a backtick answer` and `more text.`",
                Delimiter.BACKTICK,
                0,
                "this is a backtick answer",
            ),
            (
                "Here is the answer `this is a backtick answer` and `more text.`",
                Delimiter.BACKTICK,
                -1,
                "more text.",
            ),
            (
                "`Here` is the answer `this` is a `backtick answer` and `more text.`",
                Delimiter.BACKTICK,
                2,
                "backtick answer",
            ),
            (
                "Here is the answer ```this is a triple backtick answer``` and more text.",
                Delimiter.TRIPLE_BACKTICK,
                0,
                "this is a triple backtick answer",
            ),
            (
                "Here is the answer ```this``` is a ```triple``` backtick ```answer``` and more text.",
                Delimiter.TRIPLE_BACKTICK,
                -1,
                "answer",
            ),
            (
                "```Here``` is the answer ```this is a triple backtick answer``` and ```more text.```",
                Delimiter.TRIPLE_BACKTICK,
                1,
                "this is a triple backtick answer",
            ),
        ],
    )
    def test_extract_text_from_delimiters(
        self, raw_text, delimiter, match_index, expected
    ):
        result = extract_text_from_delimiters(raw_text, delimiter, match_index)
        assert result == expected

    @pytest.mark.parametrize(
        "response_text, field_list, expected",
        [
            (
                '{"name": "John", "age": "30", "city": "New York"}',
                ["name", "age", "city"],
                {"name": "John", "age": "30", "city": "New York"},
            ),
            (
                '{"name": "John", "city": "New York"}',
                ["name", "age", "city"],
                {"name": "John", "city": "New York"},
            ),
            ("", ["name", "age", "city"], {}),
            ('{"name": "John", "age": "30", "city": "New York"}', [], {}),
            (
                '{"name": "", "age": "", "city": ""}',
                ["name", "age", "city"],
                {"name": "", "age": "", "city": ""},
            ),
            (
                '{\n"name": "John",\n"age": "30",\n"city": "New York"\n}',
                ["name", "age", "city"],
                {"name": "John", "age": "30", "city": "New York"},
            ),
        ],
    )
    def test_json_parsing_search(self, response_text, field_list, expected):
        assert json_parsing_search(response_text, field_list) == expected

    @pytest.fixture
    def mock_invoke_claude(self, mocker):
        return mocker.patch.object(llm_utils, "invoke_claude")

    @pytest.mark.parametrize(
        "text, page, expected",
        [
            (
                {
                    "1": "This page contains a 10% reimbursement schedule.",
                    "2": "No relevant content here.",
                },
                "1",
                True,
            ),
            (
                {
                    "1": "This page has no relevant content.",
                    "2": "Still no relevant content here.",
                },
                "1",
                False,
            ),
            (
                {
                    "1": "This page contains a 10% reimbursement schedule.",
                    "2": "No relevant content here.",
                },
                "invalid_page",
                False,
            ),
            ("This text contains a $100 reimbursement schedule.", "1", True),
            ("This text has no relevant content.", "1", False),
        ],
    )
    def test_contains_reimbursement(self, text, page, expected):
        result = contains_reimbursement(text, page)
        assert result == expected

    def test_invalid_input_type(self, caplog):
        text = ["This is a list, not a dictionary or string."]
        page = "1"
        result = contains_reimbursement(text, page)
        assert result is False
        assert "contains_reimbursement - Invalid data type" in caplog.text

    @pytest.mark.parametrize(
        "value, expected",
        [
            (None, True),
            ("", True),
            ("N/A", True),
            ("null", True),
            ("none", True),
            ("NaN", True),
            (np.nan, True),
            (pd.NA, True),
            ("This is not empty.", False),
            (42, False),
            ("   ", True),  # Whitespace-only string
            ("\t", True),  # Tab character
            ("\n", True),  # Newline character
            ("n/a", True),  # Lowercase N/A
            ("na", True),  # Lowercase NA
            ("NULL", True),  # Uppercase NULL
            ("NONE", True),  # Uppercase NONE
            ("None", True),  # Mixed case None (already in empty_values)
            ("UNKNOWN", False),  # Should NOT be empty
            ("unknown", False),  # Should NOT be empty
            ("  N/A  ", True),  # N/A with whitespace
        ],
    )
    def test_is_empty(self, value, expected):
        result = is_empty(value)
        assert result == expected

    @pytest.mark.parametrize(
        "text_dict, filename, expected",
        [
            (
                {
                    "1": "This page has 2 signatures.",
                    "2": "None here.",
                    "3": "Signature of the party (but ignored)",
                },
                "test_file.txt",
                {"1": "This page has 2 signatures."},
            ),
            (
                {"1": "No relevant content here.", "2": "Nor over here."},
                "test_file.txt",
                {},
            ),
            (
                {
                    "1": "This page has a signature.",
                    "2": "signature page follow.",
                    "3": "Signature of the party (but ignored)",
                },
                "test_file.txt",
                {
                    "2": "signature page follow.",
                    "3": "Signature of the party (but ignored)",
                },
            ),
            (
                {
                    "1": "This page has 1 signature.",
                    "2": "No relevant content.",
                    "3": "Signature block here.  But ignored bc of the first textract-like page",
                },
                "test_file.txt",
                {"1": "This page has 1 signature."},
            ),
            (
                {"1": "No signed names at all.", "2": "Just some random text."},
                "test_file.txt",
                {},
            ),
            (
                {
                    "1": "some random text",
                    "2": "Signature authorization",
                    "3": "Signature of the party",
                },
                "test_file.txt",
                {"2": "Signature authorization", "3": "Signature of the party"},
            ),
            ({}, "test_file.txt", {}),
        ],
    )
    def test_extract_signature_page(self, text_dict, filename, expected):
        result = extract_signature_page(text_dict, filename)
        assert result == expected

    @pytest.mark.parametrize(
        "page_key, expected",
        [
            ("1", (0, 1)),  # Numeric key
            ("10", (0, 10)),  # Larger numeric key
            ("A", (1, "A")),  # Alphabetic key
            ("Z", (1, "Z")),  # Another alphabetic key
            ("1A", (1, "1A")),  # Alphanumeric key
            ("", (1, "")),  # Empty string
            ("001", (0, 1)),  # Numeric key with leading zeros
            ("B2", (1, "B2")),  # Alphanumeric key with letter first
            ("-5", (0, -5)),  # Negative numeric
            (" ", (1, " ")),  # Space character
        ],
    )
    def test_page_key_sort(self, page_key, expected):
        result = page_key_sort(page_key)
        assert result == expected

    def test_page_key_sort_with_sorting(self):
        keys = ["10", "IV", "2", "A", "1", "B", "Z"]
        expected_sorted_keys = ["1", "2", "10", "A", "B", "IV", "Z"]
        sorted_keys = sorted(keys, key=page_key_sort)
        assert sorted_keys == expected_sorted_keys


class TestCountReimbursementsInExhibit:
    @pytest.mark.parametrize(
        "exhibit_text, expected_count",
        [
            ("This exhibit includes a 10% reimbursement and a $100 reimbursement.", 2),
            ("No reimbursements mentioned here.", 0),
            ("Reimbursement of 50% and another reimbursement of $200.", 2),
            ("100 percent reimbursement and fifty dollars reimbursement.", 0),
            ("", 0),
            ("Reimbursement: 20% and $300.", 2),
            ("Reimbursement: 20% and $300. Another 10% reimbursement.", 3),
        ],
    )
    def test_count_reimbursements_in_exhibit(self, exhibit_text, expected_count):
        result = count_reimbursements_in_exhibit(exhibit_text)
        assert result == expected_count


class TestNormalizeStateToAbbreviation:
    @pytest.mark.parametrize(
        "state_value, expected",
        [
            # Special values should be preserved
            ("N/A", "N/A"),
            ("UNKNOWN", "UNKNOWN"),
            ("", ""),
            (None, None),
            # Valid 2-letter abbreviations should be uppercased
            ("ca", "CA"),
            ("CA", "CA"),
            ("ny", "NY"),
            ("TX", "TX"),
            ("fl", "FL"),
            # Full state names should be converted to abbreviations
            ("california", "CA"),
            ("California", "CA"),
            ("CALIFORNIA", "CA"),
            ("new york", "NY"),
            ("New York", "NY"),
            ("NEW YORK", "NY"),
            ("texas", "TX"),
            ("Texas", "TX"),
            ("florida", "FL"),
            ("Florida", "FL"),
            # Common variations should be handled
            ("calif", "CA"),
            ("cal", "CA"),
            ("fla", "FL"),
            ("ny", "NY"),
            ("n.y.", "NY"),
            ("tex", "TX"),
            ("penn", "PA"),
            ("mass", "MA"),
            ("wash", "WA"),
            ("washington state", "WA"),
            ("d.c.", "DC"),
            ("dc", "DC"),
            ("washington dc", "DC"),
            ("district of columbia", "DC"),
            # Edge cases with whitespace
            ("  california  ", "CA"),
            (" NY ", "NY"),
            ("	texas	", "TX"),  # tabs
            # Invalid/unrecognized values should be returned as-is
            ("invalid_state", "invalid_state"),
            ("123", "123"),
            ("XY", "XY"),  # Not a valid state abbreviation
            ("some random text", "some random text"),
            # All 50 states + DC (spot check a few more)
            ("alabama", "AL"),
            ("alaska", "AK"),
            ("arizona", "AZ"),
            ("wyoming", "WY"),
            ("hawaii", "HI"),
            ("vermont", "VT"),
        ],
    )
    def test_normalize_state_to_abbreviation(self, state_value, expected):
        from src.utils.string_utils import normalize_state_to_abbreviation

        result = normalize_state_to_abbreviation(state_value)
        assert result == expected

    def test_all_valid_abbreviations_preserved(self):
        """Test that all valid 2-letter state abbreviations are preserved when uppercase"""
        from src.utils.string_utils import normalize_state_to_abbreviation

        valid_abbreviations = [
            "AL",
            "AK",
            "AZ",
            "AR",
            "CA",
            "CO",
            "CT",
            "DE",
            "FL",
            "GA",
            "HI",
            "ID",
            "IL",
            "IN",
            "IA",
            "KS",
            "KY",
            "LA",
            "ME",
            "MD",
            "MA",
            "MI",
            "MN",
            "MS",
            "MO",
            "MT",
            "NE",
            "NV",
            "NH",
            "NJ",
            "NM",
            "NY",
            "NC",
            "ND",
            "OH",
            "OK",
            "OR",
            "PA",
            "RI",
            "SC",
            "SD",
            "TN",
            "TX",
            "UT",
            "VT",
            "VA",
            "WA",
            "WV",
            "WI",
            "WY",
            "DC",
        ]

        for abbrev in valid_abbreviations:
            result = normalize_state_to_abbreviation(abbrev)
            assert result == abbrev, f"Failed for abbreviation: {abbrev}"

            # Also test lowercase versions
            result_lower = normalize_state_to_abbreviation(abbrev.lower())
            assert (
                result_lower == abbrev
            ), f"Failed for lowercase abbreviation: {abbrev.lower()}"

    def test_case_insensitivity(self):
        """Test that the function handles various case combinations"""
        from src.utils.string_utils import normalize_state_to_abbreviation

        test_cases = [
            ("california", "CA"),
            ("California", "CA"),
            ("CALIFORNIA", "CA"),
            ("CaLiFoRnIa", "CA"),
            ("new york", "NY"),
            ("New York", "NY"),
            ("NEW YORK", "NY"),
            ("NeW yOrK", "NY"),
        ]

        for input_state, expected in test_cases:
            result = normalize_state_to_abbreviation(input_state)
            assert result == expected, f"Failed for: {input_state}"

    def test_edge_cases(self):
        """Test edge cases and boundary conditions"""
        from src.utils.string_utils import normalize_state_to_abbreviation

        # Test with just whitespace
        assert normalize_state_to_abbreviation("   ") == "   "

        # Test with mixed valid/invalid content
        assert normalize_state_to_abbreviation("california123") == "california123"

        # Test very long strings
        long_string = "this is a very long string that is not a state"
        assert normalize_state_to_abbreviation(long_string) == long_string

        # Test single characters
        assert normalize_state_to_abbreviation("C") == "C"
        assert normalize_state_to_abbreviation("A") == "A"