jacob-mathison/doczyai-pipelines
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Activity
Files
732d345d7fdae9df3ea12e2c4613e8a052b67fad
doczyai-pipelines/fieldExtraction/tests/table_utils_test.py
T
Alex Galarce d345dd0ed3 Merged in feature/integrate-standard-and-complex (pull request #356) 
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
2025-01-24 22:46:59 +00:00

730 lines
32 KiB
Python
Raw Blame History

import pytest
from src.prompts import preprocessing_prompts
import src.utils as utils
from src import config
import copy
from src.utils.table_utils import *
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 = preprocessing_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 = preprocessing_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
class TestCountTables:
def test_count_tables_single_table(self):
"""Test counting tables when there is a single table."""
exhibit_text = (
"-------Table Start-------- Table Header {'Column 1': ['Value 1', 'Value 2'], 'Column 2': ['Value 3', 'Value 4']}-------Table End--------"
)
result = count_tables(exhibit_text)
assert result == 1
def test_count_tables_multiple_tables(self):
"""Test counting tables when there are multiple tables."""
exhibit_text = (
"-------Table Start-------- Table Header {'Column 1': ['Value 1', 'Value 2'], 'Column 2': ['Value 3', 'Value 4']}-------Table End--------\n"
"-------Table Start-------- Another Table Header {'Column A': ['Value A1', 'Value A2'], 'Column B': ['Value B1', 'Value B2']}-------Table End--------"
)
result = count_tables(exhibit_text)
assert result == 2
def test_count_tables_no_table(self):
"""Test counting tables when there are no tables."""
exhibit_text = (
"This page does not contain a table.\n"
"Neither does this page."
)
result = count_tables(exhibit_text)
assert result == 0
def test_count_tables_mixed_content(self):
"""Test counting tables when there is a mix of pages with and without tables."""
exhibit_text = (
"-------Table Start-------- Table Header {'Column 1': ['Value 1', 'Value 2'], 'Column 2': ['Value 3', 'Value 4']}-------Table End--------\n"
"This page does not contain a table.\n"
"-------Table Start-------- Another Table Header {'Column A': ['Value A1', 'Value A2'], 'Column B': ['Value B1', 'Value B2']}-------Table End--------"
)
result = count_tables(exhibit_text)
assert result == 2
class TestExtractFirstTableFromPage:
def test_extract_first_table_from_page_with_table(self):
"""Test extracting the first table when the page contains a table."""
text = (
"Some text before the table.\n"
"-------Table Start--------\n"
"Table Header {'Column 1': ['Value 1', 'Value 2'], 'Column 2': ['Value 3', 'Value 4']}\n"
"-------Table End--------\n"
"Some text after the table."
)
result = extract_first_table_from_page(text)
expected = "Table Header {'Column 1': ['Value 1', 'Value 2'], 'Column 2': ['Value 3', 'Value 4']}"
assert result == expected
def test_extract_first_table_from_page_no_table(self):
"""Test extracting the first table when the page does not contain a table."""
text = "This page does not contain a table."
result = extract_first_table_from_page(text)
assert result == ""
def test_extract_first_table_from_page_multiple_tables(self):
"""Test extracting the first table when the page contains multiple tables."""
text = (
"Some text before the tables.\n"
"-------Table Start--------\n"
"First Table Header {'Column 1': ['Value 1', 'Value 2'], 'Column 2': ['Value 3', 'Value 4']}\n"
"-------Table End--------\n"
"Some text between the tables.\n"
"-------Table Start--------\n"
"Second Table Header {'Column A': ['Value A1', 'Value A2'], 'Column B': ['Value B1', 'Value B2']}\n"
"-------Table End--------\n"
"Some text after the tables."
)
result = extract_first_table_from_page(text)
expected = "First Table Header {'Column 1': ['Value 1', 'Value 2'], 'Column 2': ['Value 3', 'Value 4']}"
assert result == expected
def test_extract_first_table_from_page_no_end_marker(self):
"""Test extracting the first table when the end marker is missing."""
text = (
"Some text before the table.\n"
"-------Table Start--------\n"
"Table Header {'Column 1': ['Value 1', 'Value 2'], 'Column 2': ['Value 3', 'Value 4']}\n"
"Some text after the table."
)
result = extract_first_table_from_page(text)
assert result == ""
def test_extract_first_table_from_page_no_start_marker(self):
"""Test extracting the first table when the start marker is missing."""
text = (
"Some text before the table.\n"
"Table Header {'Column 1': ['Value 1', 'Value 2'], 'Column 2': ['Value 3', 'Value 4']}\n"
"-------Table End--------\n"
"Some text after the table."
)
result = extract_first_table_from_page(text)
assert result == ""
class TestExtractSnippetsBetweenMarkers:
def test_extract_snippets_between_markers_single_snippet(self):
"""Test extracting snippets when there is a single snippet."""
text = (
"Some text before the table.\n"
"-------Table Start--------\n"
"Table Header {'Column 1': ['Value 1', 'Value 2'], 'Column 2': ['Value 3', 'Value 4']}\n"
"-------Table End--------\n"
"Some text after the table."
)
result = extract_snippets_between_markers(text)
expected = ["\nTable Header {'Column 1': ['Value 1', 'Value 2'], 'Column 2': ['Value 3', 'Value 4']}\n"]
assert result == expected
def test_extract_snippets_between_markers_multiple_snippets(self):
"""Test extracting snippets when there are multiple snippets."""
text = (
"Some text before the tables.\n"
"-------Table Start--------\n"
"First Table Header {'Column 1': ['Value 1', 'Value 2'], 'Column 2': ['Value 3', 'Value 4']}\n"
"-------Table End--------\n"
"Some text between the tables.\n"
"-------Table Start--------\n"
"Second Table Header {'Column A': ['Value A1', 'Value A2'], 'Column B': ['Value B1', 'Value B2']}\n"
"-------Table End--------\n"
"Some text after the tables."
)
result = extract_snippets_between_markers(text)
expected = [
"\nFirst Table Header {'Column 1': ['Value 1', 'Value 2'], 'Column 2': ['Value 3', 'Value 4']}\n",
"\nSecond Table Header {'Column A': ['Value A1', 'Value A2'], 'Column B': ['Value B1', 'Value B2']}\n"
]
assert result == expected
def test_extract_snippets_between_markers_no_snippets(self):
"""Test extracting snippets when there are no snippets."""
text = "This page does not contain a table."
result = extract_snippets_between_markers(text)
assert result == []
def test_extract_snippets_between_markers_no_end_marker(self):
"""Test extracting snippets when the end marker is missing."""
text = (
"Some text before the table.\n"
"-------Table Start--------\n"
"Table Header {'Column 1': ['Value 1', 'Value 2'], 'Column 2': ['Value 3', 'Value 4']}\n"
"Some text after the table."
)
result = extract_snippets_between_markers(text)
assert result == []
def test_extract_snippets_between_markers_no_start_marker(self):
"""Test extracting snippets when the start marker is missing."""
text = (
"Some text before the table.\n"
"Table Header {'Column 1': ['Value 1', 'Value 2'], 'Column 2': ['Value 3', 'Value 4']}\n"
"-------Table End--------\n"
"Some text after the table."
)
result = extract_snippets_between_markers(text)
assert result == []
class TestGetStrDictionariesFromText:
def test_get_str_dictionaries_from_text_single_match(self):
"""Test extracting a single dictionary-like string when only_first is True."""
text = (
"Some text before the dictionary.\n"
"{'key1': 'value1', 'key2': 'value2'}\n"
"Some text after the dictionary."
)
result = get_str_dictionaries_from_text(text, only_first=True)
expected = "{'key1': 'value1', 'key2': 'value2'}"
assert result == expected
def test_get_str_dictionaries_from_text_multiple_matches(self):
"""Test extracting all dictionary-like strings when only_first is False."""
text = (
"Some text before the dictionaries.\n"
"{'key1': 'value1', 'key2': 'value2'}\n"
"Some text between the dictionaries.\n"
"{'keyA': 'valueA', 'keyB': 'valueB'}\n"
"Some text after the dictionaries."
)
result = get_str_dictionaries_from_text(text, only_first=False)
expected = [
"{'key1': 'value1', 'key2': 'value2'}",
"{'keyA': 'valueA', 'keyB': 'valueB'}"
]
assert result == expected
def test_get_str_dictionaries_from_text_no_match(self):
"""Test extracting dictionary-like strings when there are no matches."""
text = "This text does not contain any dictionary-like strings."
result = get_str_dictionaries_from_text(text, only_first=True)
assert result == ""
def test_get_str_dictionaries_from_text_empty_text(self):
"""Test extracting dictionary-like strings from an empty text."""
text = ""
result = get_str_dictionaries_from_text(text, only_first=True)
assert result == ""
def test_get_str_dictionaries_from_text_single_match_only_first_false(self):
"""Test extracting dictionary-like strings when there is a single match and only_first is False."""
text = (
"Some text before the dictionary.\n"
"{'key1': 'value1', 'key2': 'value2'}\n"
"Some text after the dictionary."
)
result = get_str_dictionaries_from_text(text, only_first=False)
expected = ["{'key1': 'value1', 'key2': 'value2'}"]
assert result == expected
class TestPageContainsMultipleTables:
def test_page_contains_multiple_tables_single_table(self):
"""Test when the page contains a single table."""
text = (
"Some text before the table.\n"
"-------Table Start--------\n"
"Table Header {'Column 1': ['Value 1', 'Value 2'], 'Column 2': ['Value 3', 'Value 4']}\n"
"-------Table End--------\n"
"Some text after the table."
)
result = page_contains_multiple_tables(text)
assert result is False
def test_page_contains_multiple_tables_multiple_tables(self):
"""Test when the page contains multiple tables."""
text = (
"Some text before the tables.\n"
"-------Table Start--------\n"
"First Table Header {'Column 1': ['Value 1', 'Value 2'], 'Column 2': ['Value 3', 'Value 4']}\n"
"-------Table End--------\n"
"Some text between the tables.\n"
"-------Table Start--------\n"
"Second Table Header {'Column A': ['Value A1', 'Value A2'], 'Column B': ['Value B1', 'Value B2']}\n"
"-------Table End--------\n"
"Some text after the tables."
)
result = page_contains_multiple_tables(text)
assert result is True
def test_page_contains_multiple_tables_no_table(self):
"""Test when the page does not contain any tables."""
text = "This page does not contain a table."
result = page_contains_multiple_tables(text)
assert result is False
def test_page_contains_multiple_tables_empty_text(self):
"""Test when the text is empty."""
text = ""
result = page_contains_multiple_tables(text)
assert result is False
def test_page_contains_multiple_tables_table_without_end_marker(self):
"""Test when the page contains a table without an end marker."""
text = (
"Some text before the table.\n"
"-------Table Start--------\n"
"Table Header {'Column 1': ['Value 1', 'Value 2'], 'Column 2': ['Value 3', 'Value 4']}\n"
"Some text after the table."
)
result = page_contains_multiple_tables(text)
assert result is False
class TestGetPageMetadata:
def test_get_page_metadata_single_table(self):
"""Test extracting metadata when the page contains a single table."""
page = (
"Page metadata before the table.\n"
"-------Table Start--------\n"
"Table Header {'Column 1': ['Value 1', 'Value 2'], 'Column 2': ['Value 3', 'Value 4']}\n"
"-------Table End--------\n"
"Some text after the table."
)
result = get_page_metadata(page)
expected = ["Page metadata before the table."]
assert result == expected
def test_get_page_metadata_multiple_tables(self):
"""Test extracting metadata when the page contains multiple tables."""
page = (
"Page metadata before the tables.\n"
"-------Table Start--------\n"
"First Table Header {'Column 1': ['Value 1', 'Value 2'], 'Column 2': ['Value 3', 'Value 4']}\n"
"-------Table End--------\n"
"Some text between the tables.\n"
"-------Table Start--------\n"
"Second Table Header {'Column A': ['Value A1', 'Value A2'], 'Column B': ['Value B1', 'Value B2']}\n"
"-------Table End--------\n"
"Some text after the tables."
)
result = get_page_metadata(page)
expected = [
("First Table Header", "{'Column 1': ['Value 1', 'Value 2'], 'Column 2': ['Value 3', 'Value 4']}"),
("Second Table Header", "{'Column A': ['Value A1', 'Value A2'], 'Column B': ['Value B1', 'Value B2']}")
]
assert result == expected
def test_get_page_metadata_no_table(self):
"""Test extracting metadata when the page does not contain any tables."""
page = "This page does not contain a table."
result = get_page_metadata(page)
expected = []
assert result == expected
def test_get_page_metadata_empty_text(self):
"""Test extracting metadata from an empty text."""
page = ""
result = get_page_metadata(page)
expected = []
assert result == expected
def test_get_page_metadata_table_without_end_marker(self):
"""Test extracting metadata when the page contains a table without an end marker."""
page = (
"Page metadata before the table.\n"
"-------Table Start--------\n"
"Table Header {'Column 1': ['Value 1', 'Value 2'], 'Column 2': ['Value 3', 'Value 4']}\n"
"Some text after the table."
)
result = get_page_metadata(page)
expected = []
assert result == expected
def test_get_page_metadata_table_without_start_marker(self):
"""Test extracting metadata when the page contains a table without a start marker."""
page = (
"Page metadata before the table.\n"
"Table Header {'Column 1': ['Value 1', 'Value 2'], 'Column 2': ['Value 3', 'Value 4']}\n"
"-------Table End--------\n"
"Some text after the table."
)
result = get_page_metadata(page)
expected = []
assert result == expected
class TestConvertStrToDict:
def test_convert_str_to_dict_valid_dict(self):
"""Test converting a valid dictionary-like string."""
text = "{'key1': 'value1', 'key2': 'value2'}"
result = convert_str_to_dict(text)
expected = {'key1': 'value1', 'key2': 'value2'}
assert result == expected
def test_convert_str_to_dict_no_dict(self):
"""Test converting a string that does not contain a dictionary."""
text = "This text does not contain any dictionary-like strings."
result = convert_str_to_dict(text)
expected = {}
assert result == expected
def test_convert_str_to_dict_empty_text(self):
"""Test converting an empty string."""
text = ""
result = convert_str_to_dict(text)
expected = {}
assert result == expected
def test_convert_str_to_dict_partial_dict(self):
"""Test converting a string that contains a partial dictionary."""
text = "Some text before the dictionary {'key1': 'value1', 'key2': 'value2'"
result = convert_str_to_dict(text)
expected = {}
assert result == expected
def test_convert_str_to_dict_multiple_dicts(self):
"""Test converting a string that contains multiple dictionaries."""
text = (
"Some text before the dictionaries.\n"
"{'key1': 'value1', 'key2': 'value2'}\n"
"Some text between the dictionaries.\n"
"{'keyA': 'valueA', 'keyB': 'valueB'}\n"
"Some text after the dictionaries."
)
result = convert_str_to_dict(text)
expected = {'key1': 'value1', 'key2': 'value2'}
assert result == expected
def test_convert_str_to_dict_invalid_dict(self):
"""Test converting a string that contains an invalid dictionary."""
text = "{'key1': 'value1', 'key2': value2}" # Missing quotes around value2
result = convert_str_to_dict(text)
expected = {}
assert result == expected
class TestInsertMetadata:
def test_insert_metadata_with_valid_data(self):
"""Test inserting metadata with valid data."""
metadata = "Page metadata"
table_content = "Table Header {'Column 1': ['Value 1', 'Value 2'], 'Column 2': ['Value 3', 'Value 4']}"
result = insert_metadata(metadata, table_content)
expected = "Page metadata\nTable Header {'Column 1': ['Value 1', 'Value 2'], 'Column 2': ['Value 3', 'Value 4']}"
assert result == expected
def test_insert_metadata_with_empty_metadata(self):
"""Test inserting metadata when metadata is empty."""
metadata = ""
table_content = "Table Header {'Column 1': ['Value 1', 'Value 2'], 'Column 2': ['Value 3', 'Value 4']}"
result = insert_metadata(metadata, table_content)
expected = "\nTable Header {'Column 1': ['Value 1', 'Value 2'], 'Column 2': ['Value 3', 'Value 4']}"
assert result == expected
def test_insert_metadata_with_empty_table_content(self):
"""Test inserting metadata when table content is empty."""
metadata = "Page metadata"
table_content = ""
result = insert_metadata(metadata, table_content)
expected = "Page metadata\n"
assert result == expected
def test_insert_metadata_with_both_empty(self):
"""Test inserting metadata when both metadata and table content are empty."""
metadata = ""
table_content = ""
result = insert_metadata(metadata, table_content)
expected = "\n"
assert result == expected
def test_insert_metadata_with_special_characters(self):
"""Test inserting metadata when metadata and table content contain special characters."""
metadata = "Page metadata with special characters: !@#$%^&*()"
table_content = "Table Header {'Column 1': ['Value 1', 'Value 2'], 'Column 2': ['Value 3', 'Value 4']} with special characters: <>?/\\|"
result = insert_metadata(metadata, table_content)
expected = "Page metadata with special characters: !@#$%^&*()\nTable Header {'Column 1': ['Value 1', 'Value 2'], 'Column 2': ['Value 3', 'Value 4']} with special characters: <>?/\\|"
assert result == expected
class TestInsertColumnHeaders:
def test_insert_column_headers_with_valid_data(self):
"""Test inserting column headers with valid data."""
data_dict = {
"Column 1": ["Value 1", "Value 2"],
"Column 2": ["Value 3", "Value 4"]
}
columns = ["Column 1", "Column 2"]
result = insert_column_headers(data_dict, columns)
expected = {
"Column 1": ["Column 1", "Value 1", "Value 2"],
"Column 2": ["Column 2", "Value 3", "Value 4"]
}
assert result == expected
def test_insert_column_headers_with_empty_data_dict(self):
"""Test inserting column headers when data_dict is empty."""
data_dict = {}
columns = ["Column 1", "Column 2"]
result = insert_column_headers(data_dict, columns)
expected = {}
assert result == expected
def test_insert_column_headers_with_mismatched_columns(self):
"""Test inserting column headers when columns list does not match data_dict keys."""
data_dict = {
"Column 1": ["Value 1", "Value 2"],
"Column 2": ["Value 3", "Value 4"]
}
columns = ["Column A", "Column B"]
result = insert_column_headers(data_dict, columns)
expected = {
"Column A": ["Column 1", "Value 1", "Value 2"],
"Column B": ["Column 2", "Value 3", "Value 4"]
}
assert result == expected
def test_insert_column_headers_with_extra_columns(self):
"""Test inserting column headers when columns list has extra headers."""
data_dict = {
"Column 1": ["Value 1", "Value 2"],
"Column 2": ["Value 3", "Value 4"]
}
columns = ["Column 1", "Column 2", "Column 3"]
result = insert_column_headers(data_dict, columns)
expected = {
"Column 1": ["Column 1", "Value 1", "Value 2"],
"Column 2": ["Column 2", "Value 3", "Value 4"]
}
assert result == expected
def test_insert_column_headers_with_missing_columns(self):
"""Test inserting column headers when columns list has missing headers."""
data_dict = {
"Column 1": ["Value 1", "Value 2"],
"Column 2": ["Value 3", "Value 4"]
}
columns = ["Column 1"]
result = insert_column_headers(data_dict, columns)
expected = {
"Column 1": ["Column 1", "Value 1", "Value 2"]
}
assert result == expected
def test_insert_column_headers_with_empty_columns(self):
"""Test inserting column headers when columns list is empty."""
data_dict = {
"Column 1": ["Value 1", "Value 2"],
"Column 2": ["Value 3", "Value 4"]
}
columns = []
result = insert_column_headers(data_dict, columns)
expected = {}
assert result == expected
class TestSplitTablesToSeparateSubpages:
def test_split_tables_to_separate_subpages_single_table(self):
"""Test splitting pages with a single table."""
text_dict = {
"1": "Some text before the table.\n"
"-------Table Start--------\n"
"Table Header {'Column 1': ['Value 1', 'Value 2'], 'Column 2': ['Value 3', 'Value 4']}\n"
"-------Table End--------\n"
"Some text after the table."
}
result = split_tables_to_separate_subpages(text_dict)
expected = {
"1.0": "Some text before the table.\n"
"-------Table Start--------\n"
"Table Header {'Column 1': ['Value 1', 'Value 2'], 'Column 2': ['Value 3', 'Value 4']}\n"
"-------Table End--------",
"1.1": "\nSome text after the table."
}
assert result == expected
def test_split_tables_to_separate_subpages_multiple_tables(self):
"""Test splitting pages with multiple tables."""
text_dict = {
"1": "Some text before the tables.\n"
"-------Table Start--------\n"
"First Table Header {'Column 1': ['Value 1', 'Value 2'], 'Column 2': ['Value 3', 'Value 4']}\n"
"-------Table End--------\n"
"Some text between the tables.\n"
"-------Table Start--------\n"
"Second Table Header {'Column A': ['Value A1', 'Value A2'], 'Column B': ['Value B1', 'Value B2']}\n"
"-------Table End--------\n"
"Some text after the tables."
}
result = split_tables_to_separate_subpages(text_dict)
expected = {
"1.0": "Some text before the tables.\n"
"-------Table Start--------\n"
"First Table Header {'Column 1': ['Value 1', 'Value 2'], 'Column 2': ['Value 3', 'Value 4']}\n"
"-------Table End--------",
"1.1": "\nSome text between the tables.\n"
"-------Table Start--------\n"
"Second Table Header {'Column A': ['Value A1', 'Value A2'], 'Column B': ['Value B1', 'Value B2']}\n"
"-------Table End--------",
"1.2": "\nSome text after the tables."
}
assert result == expected
def test_split_tables_to_separate_subpages_no_table(self):
"""Test splitting pages with no tables."""
text_dict = {
"1": "This page does not contain a table."
}
result = split_tables_to_separate_subpages(text_dict)
expected = {
"1": "This page does not contain a table."
}
assert result == expected
def test_split_tables_to_separate_subpages_empty_text(self):
"""Test splitting pages with empty text."""
text_dict = {
"1": ""
}
result = split_tables_to_separate_subpages(text_dict)
expected = {}
assert result == expected
def test_split_tables_to_separate_subpages_mixed_content(self):
"""Test splitting pages with mixed content."""
text_dict = {
"1": "Some text before the table.\n"
"-------Table Start--------\n"
"Table Header {'Column 1': ['Value 1', 'Value 2'], 'Column 2': ['Value 3', 'Value 4']}\n"
"-------Table End--------\n"
"Some text after the table.",
"2": "This page does not contain a table.",
"3": "Some text before the tables.\n"
"-------Table Start--------\n"
"First Table Header {'Column 1': ['Value 1', 'Value 2'], 'Column 2': ['Value 3', 'Value 4']}\n"
"-------Table End--------\n"
"Some text between the tables.\n"
"-------Table Start--------\n"
"Second Table Header {'Column A': ['Value A1', 'Value A2'], 'Column B': ['Value B1', 'Value B2']}\n"
"-------Table End--------\n"
"Some text after the tables."
}
result = split_tables_to_separate_subpages(text_dict)
expected = {
"1.0": "Some text before the table.\n"
"-------Table Start--------\n"
"Table Header {'Column 1': ['Value 1', 'Value 2'], 'Column 2': ['Value 3', 'Value 4']}\n"
"-------Table End--------",
"1.1": "\nSome text after the table.",
"2": "This page does not contain a table.",
"3.0": "Some text before the tables.\n"
"-------Table Start--------\n"
"First Table Header {'Column 1': ['Value 1', 'Value 2'], 'Column 2': ['Value 3', 'Value 4']}\n"
"-------Table End--------",
"3.1": "\nSome text between the tables.\n"
"-------Table Start--------\n"
"Second Table Header {'Column A': ['Value A1', 'Value A2'], 'Column B': ['Value B1', 'Value B2']}\n"
"-------Table End--------",
"3.2": "\nSome text after the tables."
}
assert result == expected
Reference in New Issue View Git Blame Copy Permalink