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

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