Siddhant Medar
ab71bd8dae
Task/unit tests investment and client * move io_utils.py * add __init__.py * refactor: move string_funcs and update import paths to use utils namespace * rename string_funcs to string_utils * move and rename table_utils.py * refactor: rename claude_funcs to llm_utils in multiple files * update unit tests * added unit tests for src/string_utils.py * added unit tests for src/table_utils.py * added unit tests for src/llm_utils.py * added unit tests for src/io_utils.py * Merge branch 'main' into task/unit-tests-investment-and-client; TODO: fix unit tests * Unit tests to fix * fix unit tests * fix invalid references Approved-by: Alex Galarce
114 lines
4.8 KiB
Python
114 lines
4.8 KiB
Python
import pytest
|
|
import src.utils as utils
|
|
from src.utils.table_utils import align_and_format_tables
|
|
from src import (prompts, config)
|
|
import copy
|
|
|
|
class TestAlignAndFormatTables:
|
|
@pytest.fixture
|
|
def mock_utils(self, mocker):
|
|
mock_contains_reimbursement = mocker.patch.object(
|
|
utils.string_utils,
|
|
"contains_reimbursement",
|
|
)
|
|
mock_invoke_claude = mocker.patch.object(
|
|
utils.llm_utils,
|
|
"invoke_claude",
|
|
)
|
|
return mock_contains_reimbursement, mock_invoke_claude
|
|
|
|
def test_align_and_format_tables_with_reimbursement(self, mock_utils):
|
|
"""Test aligning and formatting tables when reimbursement keywords are present."""
|
|
mock_contains_reimbursement, mock_invoke_claude = mock_utils
|
|
mock_contains_reimbursement.return_value = True
|
|
mock_invoke_claude.return_value = (
|
|
"Table Header\nColumn 1 : Value 1% | Column 2 : Value 3\nColumn 1 : Value 2 | Column 2 : Value 4"
|
|
)
|
|
|
|
text_dict = {
|
|
"1": "-------Table Start-------- Table Header {'Column 1': ['Value 1%', 'Value 2'], 'Column 2': ['Value 3', 'Value 4']}-------Table End--------",
|
|
}
|
|
filename = "test_file.txt"
|
|
|
|
expected_prompt = prompts.ALIGN_AND_FORMAT_TABLES(text_dict["1"])
|
|
input_text_dict = copy.deepcopy(text_dict)
|
|
|
|
result = align_and_format_tables(text_dict, filename)
|
|
|
|
assert result == {
|
|
"1": "Table Header\nColumn 1 : Value 1% | Column 2 : Value 3\nColumn 1 : Value 2 | Column 2 : Value 4",
|
|
}
|
|
mock_contains_reimbursement.assert_called_once_with(input_text_dict["1"])
|
|
mock_invoke_claude.assert_called_once_with(
|
|
expected_prompt,
|
|
config.MODEL_ID_CLAUDE35_SONNET,
|
|
filename,
|
|
8192,
|
|
)
|
|
|
|
def test_align_and_format_tables_without_reimbursement(self, mock_utils):
|
|
"""Test aligning and formatting tables when no reimbursement keywords are present."""
|
|
mock_contains_reimbursement, mock_invoke_claude = mock_utils
|
|
mock_contains_reimbursement.return_value = False
|
|
|
|
text_dict = {
|
|
"1": "-------Table Start-------- Table Header {'Column 1': ['Value 1', 'Value 2'], 'Column 2': ['Value 3', 'Value 4']}-------Table End--------",
|
|
}
|
|
filename = "test_file.txt"
|
|
|
|
result = align_and_format_tables(text_dict, filename)
|
|
|
|
assert result == text_dict # No changes should be made
|
|
mock_contains_reimbursement.assert_called_once_with(text_dict["1"])
|
|
mock_invoke_claude.assert_not_called()
|
|
|
|
def test_align_and_format_tables_no_table(self, mock_utils):
|
|
"""Test aligning and formatting tables when no table is present."""
|
|
mock_contains_reimbursement, mock_invoke_claude = mock_utils
|
|
|
|
text_dict = {
|
|
"1": "This page does not contain a table.",
|
|
}
|
|
filename = "test_file.txt"
|
|
|
|
result = align_and_format_tables(text_dict, filename)
|
|
|
|
assert result == text_dict # No changes should be made
|
|
mock_contains_reimbursement.assert_not_called()
|
|
mock_invoke_claude.assert_not_called()
|
|
|
|
def test_align_and_format_tables_multiple_pages(self, mock_utils):
|
|
"""Test aligning and formatting tables with multiple pages."""
|
|
mock_contains_reimbursement, mock_invoke_claude = mock_utils
|
|
mock_contains_reimbursement.side_effect = [True, False]
|
|
mock_invoke_claude.return_value = (
|
|
"Table Header\nColumn 1 : Value 1 | Column 2 : Value 3%\nColumn 1 : Value 2 | Column 2 : Value 4"
|
|
)
|
|
|
|
text_dict = {
|
|
"1": "-------Table Start-------- Table Header {'Column 1': ['Value 1', 'Value 2'], 'Column 2': ['Value 3%', 'Value 4']}-------Table End--------",
|
|
"2": "This page has 'Table Start' but does not contain reimbursement.",
|
|
}
|
|
filename = "test_file.txt"
|
|
|
|
expected_prompt = prompts.ALIGN_AND_FORMAT_TABLES(text_dict["1"])
|
|
input_text_dict = copy.deepcopy(text_dict)
|
|
|
|
result = align_and_format_tables(text_dict, filename)
|
|
|
|
assert result == {
|
|
"1": "Table Header\nColumn 1 : Value 1 | Column 2 : Value 3%\nColumn 1 : Value 2 | Column 2 : Value 4",
|
|
"2": "This page has 'Table Start' but does not contain reimbursement.",
|
|
}
|
|
assert mock_contains_reimbursement.call_count == 2
|
|
mock_contains_reimbursement.assert_any_call(input_text_dict["1"])
|
|
mock_contains_reimbursement.assert_any_call(input_text_dict["2"])
|
|
mock_invoke_claude.assert_called_once_with(
|
|
expected_prompt,
|
|
config.MODEL_ID_CLAUDE35_SONNET,
|
|
filename,
|
|
8192,
|
|
)
|
|
for call in mock_invoke_claude.call_args_list:
|
|
prompt, _, _, _ = call[0] # Unpack the arguments
|
|
assert input_text_dict["2"] not in prompt |