Alex Galarce
d345dd0ed3
Feature/integrate standard and complex * last edits to merge_tables_draft_v2 before moving to table_funcs * migrate to table_utils.py * fix typehint that caused mypy error * fix mypy errors * black and isort * fix error with multi-table pages * update poetry.lock * add tests, fix bug when table start marker or table end marker is missing * added more tests * tests for get_str_dictionaries_from_text() * fix edge case in get_str_dictionaries_from_text * fix mypy error * more tests * add tests, add handling for invalid dictionaries * add tests * add tests for insert_column_headers() * update requirements * add new simple/complex logic to preprocess.py * remove files used solely for testing * split pages into sub-pages by tables * add table split on end marker. Also add tests * add docstrings, change control flow * remove unused tests, add TODO for new tests * get in prompt changes and intermediate decisions * TODO for future enhancement for get_exhibit_pages Approved-by: Katon Minhas
211 lines
9.6 KiB
Python
211 lines
9.6 KiB
Python
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
|