doczyai-pipelines/fieldExtraction/tests/string_utils_test.py

import pytest
import src.utils as utils
from src.utils.string_utils import extract_text_from_delimiters, json_parsing_search, secondary_string_to_dict, primary_string_to_dict, contains_reimbursement, is_empty, count_reimbursements_in_exhibit, extract_signature_page, page_key_sort
import src.utils.llm_utils as llm_utils
import numpy as np
import pandas as pd

from src.enums.delimiters import Delimiter

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')

    def test_valid_json(self, mock_invoke_claude):
        dict_string = '{"key": "value"}'
        filename = "test_file.txt"
        result = secondary_string_to_dict(dict_string, filename)
        assert result == {"key": "value"}
        mock_invoke_claude.assert_not_called()

    def test_invalid_json_fixed_by_replacing_quotes(self, mock_invoke_claude):
        dict_string = "{'key': 'value'}"
        filename = "test_file.txt"
        result = secondary_string_to_dict(dict_string, filename)
        assert result == {"key": "value"}
        mock_invoke_claude.assert_not_called()

    def test_invalid_json_fixed_by_llm(self, mock_invoke_claude):
        mock_invoke_claude.return_value = '{"key": "value"}'
        dict_string = "{key: value}"
        filename = "test_file.txt"
        result = secondary_string_to_dict(dict_string, filename)
        assert result == {"key": "value"}
        mock_invoke_claude.assert_called_once()

    def test_invalid_json_failed_by_llm(self, mock_invoke_claude):
        mock_invoke_claude.side_effect = Exception("LLM call failed")
        dict_string = "{key: value}"
        filename = "test_file.txt"
        result = secondary_string_to_dict(dict_string, filename)
        assert result == {}
        mock_invoke_claude.assert_called_once()

    def test_no_dict_found(self, mock_invoke_claude):
        mock_invoke_claude.return_value = '{}'
        dict_string = "This is a plain text without any dictionary."
        filename = "test_file.txt"
        result = secondary_string_to_dict(dict_string, filename)
        assert result == {}
        mock_invoke_claude.assert_called_once()

    @pytest.fixture
    def mock_secondary_string_to_dict(self, mocker):
        return mocker.patch.object(utils.string_utils, "secondary_string_to_dict")

    def test_valid_input(self, mock_secondary_string_to_dict):
        mock_secondary_string_to_dict.side_effect = [
            {"key1": "value1"},
            {"key2": "value2"},
            {"key3": "value3"},
        ]
        string_dict = {
            "1": "[{'key1': 'value1'}, {'key2': 'value2'}]",
            "2": "[{'key3': 'value3'}]",
        }
        filename = "test_file.txt"
        expected_output = [
            {"key1": "value1", "page_num": "1"},
            {"key2": "value2", "page_num": "1"},
            {"key3": "value3", "page_num": "2"},
        ]
        result = primary_string_to_dict(string_dict, filename)
        assert result == expected_output
        assert mock_secondary_string_to_dict.call_count == len(expected_output)

    def test_invalid_input_no_dicts(self, mock_secondary_string_to_dict):
        mock_secondary_string_to_dict.return_value = {}
        string_dict = {
            "1": "[This is not a dictionary]",
            "2": "[Still not a dictionary]",
        }
        filename = "test_file.txt"
        expected_output = []
        result = primary_string_to_dict(string_dict, filename)
        assert result == expected_output
        assert mock_secondary_string_to_dict.call_count == len(expected_output)

    def test_empty_input(self, mock_secondary_string_to_dict):
        string_dict = {}
        filename = "test_file.txt"
        expected_output = []
        result = primary_string_to_dict(string_dict, filename)
        assert result == expected_output
        mock_secondary_string_to_dict.assert_not_called()

    def test_malformed_input(self, mock_secondary_string_to_dict):
        mock_secondary_string_to_dict.side_effect = [
            {"key1": "value1"},
            {"key2": "value2"},
        ]
        string_dict = {
            "1": "[{'key1': 'value1'}, malformed]",
            "2": "[{'key2': 'value2'}, <<<invalid>>>]",
        }
        filename = "test_file.txt"
        expected_output = [
            {"key1": "value1", "page_num": "1"},
            {"key2": "value2", "page_num": "2"},
        ]
        result = primary_string_to_dict(string_dict, filename)
        assert result == expected_output
        assert mock_secondary_string_to_dict.call_count == 2

    def test_multiple_pages(self, mock_secondary_string_to_dict):
        mock_secondary_string_to_dict.side_effect = [
            {"key1": "value1"},
            {"key2": "value2"},
            {"key3": "value3"},
            {"key4": "value4"},
        ]
        string_dict = {
            "1": "[{'key1': 'value1'}, {'key2': 'value2'}]",
            "2": "[{'key3': 'value3'}, {'key4': 'value4'}]",
        }
        filename = "test_file.txt"
        expected_output = [
            {"key1": "value1", "page_num": "1"},
            {"key2": "value2", "page_num": "1"},
            {"key3": "value3", "page_num": "2"},
            {"key4": "value4", "page_num": "2"},
        ]
        result = primary_string_to_dict(string_dict, filename)
        assert result == expected_output
        assert mock_secondary_string_to_dict.call_count == 4

    @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, capsys):
        text = ["This is a list, not a dictionary or string."]
        page = "1"
        result = contains_reimbursement(text, page)
        assert result is False
        captured = capsys.readouterr()
        assert "contains_reimbursement - Invalid data type" in captured.out

    @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),
    ])
    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": "Another signature is found here."},
                "test_file.txt",
                {"1": "This page has a signature.", "2": "Another signature is found here."}
            ),
            (
                {"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": "This page has a signature.", "2": "Signature is mentioned here too."},
                "test_file.txt",
                {"1": "This page has a signature.", "2": "Signature is mentioned here too."}
            ),
            (
                {},
                "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