doczyai-pipelines/fieldExtraction/tests/table_utils_test.py

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