Alex Galarce
674add9a3e
Feature/black and isort * black and isort constants * black and isort codes * black and isort crosswalks * black and isort investment * black and isort prompts * isort testbed * isort tracking * black and isort utils * poetry and black tests Approved-by: Katon Minhas
1081 lines
43 KiB
Python
1081 lines
43 KiB
Python
import logging
|
|
import unittest
|
|
from unittest.mock import patch
|
|
|
|
import pandas as pd
|
|
from constants.constants import Constants
|
|
from src.investment.table_funcs import (END_MARKER, START_MARKER, Table,
|
|
combine_continuous_tables,
|
|
convert_str_to_dataframe,
|
|
get_table_info, recreate_page,
|
|
split_large_table_with_headers,
|
|
split_tables_by_rows)
|
|
|
|
logging.basicConfig(level=logging.INFO)
|
|
|
|
|
|
class TestTableFunctions(unittest.TestCase):
|
|
def setUp(self):
|
|
# New list-of-rows format
|
|
self.sample_text_dict_new = {
|
|
"1": "Some metadata\n-------Table Start--------\nTable 1\n[['A', 'B'], ['1%', '3%'], ['2', '4']]\n-------Table End--------",
|
|
"2": "More metadata\n-------Table Start--------\n\n[['100', '200'], ['5$', '7$'], ['6', '8']]\n-------Table End--------",
|
|
}
|
|
# Old dict format for backward compatibility testing
|
|
self.sample_text_dict_old = {
|
|
"1": "Some metadata\n-------Table Start--------\nTable 1\n{'A': ['1', '2$'], 'B': ['3$', '4']}\n-------Table End--------",
|
|
"2": "More metadata\n-------Table Start--------\n\n{'A': ['5', '6%'], 'B': ['7%', '8']}\n-------Table End--------",
|
|
}
|
|
# Multiple tables on one page
|
|
self.multi_table_text = {
|
|
"1": "Page metadata\n-------Table Start--------\nFirst Table\n[['Col1', 'Col2'], ['A%', 'B'], ['C', 'D']]\n-------Table End--------\n\nSome text\n-------Table Start--------\nSecond Table\n[['X', 'Y'], ['1', '2']]\n-------Table End--------"
|
|
}
|
|
# Long table for splitting tests
|
|
self.long_table_text = {
|
|
"1": "Long table metadata\n-------Table Start--------\nLong Table\n[['Name%', 'Value'], ['Row1', '1'], ['Row2', '2'], ['Row3', '3'], ['Row4', '4'], ['Row5', '5']]\n-------Table End--------"
|
|
}
|
|
|
|
self.EXHIBIT_HEADER_MARKERS = Constants().EXHIBIT_HEADER_MARKERS
|
|
|
|
def test_convert_str_to_dataframe_new_format(self):
|
|
"""Test conversion from new list-of-rows format."""
|
|
table_str = "[['A', 'B'], ['1', '3'], ['2', '4']]"
|
|
df = convert_str_to_dataframe(table_str)
|
|
|
|
self.assertIsInstance(df, pd.DataFrame)
|
|
self.assertEqual(df.shape, (2, 2)) # 2 rows, 2 columns
|
|
self.assertEqual(df.columns[0], "A")
|
|
self.assertEqual(df.iloc[0, 0], "1")
|
|
self.assertEqual(df.iloc[1, 0], "2")
|
|
self.assertEqual(df.iloc[0, 1], "3")
|
|
self.assertTrue(
|
|
df.dtypes.eq("object").all()
|
|
) # All columns should be of type str (object)
|
|
|
|
def test_convert_str_to_dataframe_old_format(self):
|
|
"""Test conversion from old dict format with deprecation warning."""
|
|
table_str = "{'A': ['1', '2'], 'B': ['3', '4']}"
|
|
with self.assertLogs(level="INFO") as log:
|
|
df = convert_str_to_dataframe(table_str)
|
|
|
|
# Should warn about old format
|
|
self.assertTrue(
|
|
any("old table (dict) format detected" in message for message in log.output)
|
|
)
|
|
|
|
self.assertIsInstance(df, pd.DataFrame)
|
|
self.assertEqual(df.shape, (2, 2)) # 2 rows, 2 columns
|
|
self.assertTrue(
|
|
df.dtypes.eq("object").all()
|
|
) # All columns should be of type str (object)
|
|
self.assertEqual(df.iloc[0, 0], "1")
|
|
self.assertEqual(df.iloc[0, 1], "3")
|
|
self.assertEqual(df.iloc[1, 0], "2")
|
|
self.assertEqual(df.iloc[1, 1], "4")
|
|
|
|
def test_convert_str_to_dataframe_invalid_format(self):
|
|
"""Test handling of invalid formats."""
|
|
invalid_str = "not a valid table format"
|
|
df = convert_str_to_dataframe(invalid_str)
|
|
|
|
self.assertIsInstance(df, pd.DataFrame)
|
|
self.assertTrue(df.empty)
|
|
|
|
def test_table_init_new_format(self):
|
|
"""Test Table initialization with new list-of-rows format."""
|
|
table_data = [["A", "B"], ["1", "3"], ["2", "4"]]
|
|
table = Table(
|
|
page_num="1",
|
|
metadata="metadata",
|
|
header="Table 1",
|
|
table_data=table_data,
|
|
start_index=0,
|
|
end_index=100,
|
|
)
|
|
|
|
self.assertEqual(table.page_num, "1")
|
|
self.assertEqual(table.metadata, "metadata")
|
|
self.assertEqual(table.header, "Table 1")
|
|
self.assertEqual(table.table.shape, (2, 2)) # 2 rows, 2 columns
|
|
self.assertEqual(list(table.table.columns), ["A", "B"])
|
|
self.assertEqual(table.table.iloc[0, 0], "1")
|
|
|
|
def test_table_init_dataframe(self):
|
|
"""Test Table initialization with a DataFrame."""
|
|
table_data = pd.DataFrame([["A", "B"], ["1", "3"]], dtype=str)
|
|
table = Table(
|
|
page_num="1",
|
|
metadata="metadata",
|
|
header="Table 1",
|
|
table_data=table_data,
|
|
start_index=0,
|
|
end_index=100,
|
|
)
|
|
self.assertEqual(table.table.shape, (2, 2)) # 2 rows, 2 columns
|
|
self.assertTrue(
|
|
table.table.dtypes.eq("object").all()
|
|
) # All columns should be of type str (object)
|
|
|
|
def test_table_init_column_mismatch_padding(self):
|
|
"""Test table initialization handles column count mismatches by padding."""
|
|
# Headers have 3 columns, but data rows have different counts
|
|
table_data = [
|
|
["A", "B", "C"], # 3 columns (header)
|
|
["1", "2"], # 2 columns - should be padded
|
|
["3", "4", "5", "6"], # 4 columns - determines max_cols = 4
|
|
["7"], # 1 column - should be padded
|
|
]
|
|
table = Table(
|
|
page_num="1",
|
|
metadata="metadata",
|
|
header="Test Table",
|
|
table_data=table_data,
|
|
start_index=0,
|
|
end_index=100,
|
|
)
|
|
|
|
# Should have 4 columns (max from all rows) and 3 data rows
|
|
self.assertEqual(table.table.shape, (3, 4))
|
|
self.assertEqual(
|
|
list(table.table.columns), ["A", "B", "C", ""]
|
|
) # Header padded too
|
|
|
|
# Check first row (padded to 4 columns)
|
|
self.assertEqual(table.table.iloc[0, 0], "1")
|
|
self.assertEqual(table.table.iloc[0, 1], "2")
|
|
self.assertEqual(table.table.iloc[0, 2], "") # Padded
|
|
self.assertEqual(table.table.iloc[0, 3], "") # Padded
|
|
|
|
# Check second row (kept at 4 columns)
|
|
self.assertEqual(table.table.iloc[1, 0], "3")
|
|
self.assertEqual(table.table.iloc[1, 1], "4")
|
|
self.assertEqual(table.table.iloc[1, 2], "5")
|
|
self.assertEqual(table.table.iloc[1, 3], "6") # 4th column preserved
|
|
|
|
# Check third row (heavily padded)
|
|
self.assertEqual(table.table.iloc[2, 0], "7")
|
|
self.assertEqual(table.table.iloc[2, 1], "") # Padded
|
|
self.assertEqual(table.table.iloc[2, 2], "") # Padded
|
|
self.assertEqual(table.table.iloc[2, 3], "") # Padded
|
|
|
|
def test_table_init_continuation_column_handling(self):
|
|
"""Test continuation table handles column mismatches with generic column names."""
|
|
# Continuation table - no header row, just data with varying lengths
|
|
table_data = [
|
|
["1", "2"], # 2 columns
|
|
["3", "4", "5"], # 3 columns
|
|
["6"], # 1 column
|
|
]
|
|
table = Table(
|
|
page_num="2",
|
|
metadata="",
|
|
header="",
|
|
table_data=table_data,
|
|
start_index=0,
|
|
end_index=100,
|
|
is_continuation=True,
|
|
)
|
|
|
|
# Should use max column count (3) and generate generic column names
|
|
self.assertEqual(table.table.shape, (3, 3))
|
|
self.assertEqual(list(table.table.columns), ["col_0", "col_1", "col_2"])
|
|
|
|
# Check padding/truncation
|
|
self.assertEqual(table.table.iloc[0, 0], "1")
|
|
self.assertEqual(table.table.iloc[0, 1], "2")
|
|
self.assertEqual(table.table.iloc[0, 2], "") # Padded
|
|
|
|
self.assertEqual(table.table.iloc[1, 0], "3")
|
|
self.assertEqual(table.table.iloc[1, 1], "4")
|
|
self.assertEqual(table.table.iloc[1, 2], "5")
|
|
|
|
self.assertEqual(table.table.iloc[2, 0], "6")
|
|
self.assertEqual(table.table.iloc[2, 1], "") # Padded
|
|
self.assertEqual(table.table.iloc[2, 2], "") # Padded
|
|
|
|
def test_combine_continuous_tables_simple(self):
|
|
"""Test the simplified continuous table combination logic."""
|
|
# Create test data with continuation table
|
|
text_dict = {
|
|
"1": "Metadata 1\n-------Table Start--------\nMain Table\n[['Name', 'Value'], ['A', '1%'], ['B', '2']]\n-------Table End--------",
|
|
"2": "-------Table Start--------\n\n[['C', '3$'], ['D', '4']]\n-------Table End--------", # Continuation (empty header)
|
|
"3": "Metadata 3\n-------Table Start--------\nAnother Table\n[['X', 'Y'], ['Z', '5%']]\n-------Table End--------",
|
|
}
|
|
|
|
# Debug: Let's see what get_table_info thinks about these tables
|
|
table_dict = get_table_info(text_dict, self.EXHIBIT_HEADER_MARKERS)
|
|
print(f"Table dict keys: {list(table_dict.keys())}")
|
|
for page_num, tables in table_dict.items():
|
|
for i, table in enumerate(tables):
|
|
print(
|
|
f"Page {page_num}, Table {i}: is_continuation={table.is_continuation}, header='{table.header}'"
|
|
)
|
|
|
|
result = combine_continuous_tables(text_dict, table_dict)
|
|
print(f"Result keys: {list(result.keys())}")
|
|
|
|
# Page 2 should be removed (combined into page 1)
|
|
self.assertNotIn("2", result)
|
|
self.assertIn("1", result)
|
|
self.assertIn("3", result)
|
|
|
|
# Page 1 should now contain the combined table
|
|
combined_table_dict = get_table_info(
|
|
{"1": result["1"]}, self.EXHIBIT_HEADER_MARKERS
|
|
)
|
|
combined_table = combined_table_dict["1"][0]
|
|
|
|
# Should have 4 rows total (2 from main + 2 from continuation)
|
|
self.assertEqual(combined_table.table.shape, (4, 2))
|
|
self.assertEqual(combined_table.table.iloc[0, 0], "A") # First table data
|
|
self.assertEqual(combined_table.table.iloc[2, 0], "C") # Continuation data
|
|
self.assertEqual(combined_table.table.iloc[3, 0], "D") # Continuation data
|
|
|
|
def test_combine_continuous_tables_multiple_tables_per_page(self):
|
|
"""Test that continuation tables combine with the LAST table on previous page, not the first."""
|
|
# Page 1 has TWO tables: Table A (complete) and Table B (continues on page 2)
|
|
text_dict = {
|
|
"1": (
|
|
"Metadata 1\n"
|
|
"-------Table Start--------\nTable A (Complete)\n[['Col1', 'Col2'], ['A1%', 'A2'], ['A3', 'A4']]\n-------Table End--------\n"
|
|
"Some text between tables\n"
|
|
"-------Table Start--------\nTable B (Continues)\n[['Name', 'Value'], ['B1%', '1'], ['B2', '2']]\n-------Table End--------"
|
|
),
|
|
"2": (
|
|
"-------Table Start--------\n\n[['B3', '3'], ['B4', '4%']]\n-------Table End--------" # Continuation of Table B
|
|
),
|
|
}
|
|
|
|
table_dict = get_table_info(text_dict, self.EXHIBIT_HEADER_MARKERS)
|
|
|
|
# Verify setup: Page 1 should have 2 tables, Page 2 should have 1 continuation table
|
|
self.assertEqual(len(table_dict["1"]), 2)
|
|
self.assertEqual(len(table_dict["2"]), 1)
|
|
self.assertFalse(
|
|
table_dict["1"][0].is_continuation
|
|
) # Table A - not continuation
|
|
self.assertFalse(
|
|
table_dict["1"][1].is_continuation
|
|
) # Table B - not continuation
|
|
self.assertTrue(
|
|
table_dict["2"][0].is_continuation
|
|
) # Table B continuation - is continuation
|
|
|
|
# Get original row counts before combining
|
|
table_a_original_rows = table_dict["1"][0].table.shape[0] # Should be 2 rows
|
|
table_b_original_rows = table_dict["1"][1].table.shape[0] # Should be 2 rows
|
|
continuation_rows = table_dict["2"][0].table.shape[0] # Should be 2 rows
|
|
|
|
# Run the combination
|
|
result = combine_continuous_tables(text_dict, table_dict)
|
|
|
|
# Page 2 should be removed (combined into page 1)
|
|
self.assertNotIn("2", result)
|
|
self.assertIn("1", result)
|
|
|
|
# Parse the combined result to verify tables
|
|
combined_table_dict = get_table_info(
|
|
{"1": result["1"]}, self.EXHIBIT_HEADER_MARKERS
|
|
)
|
|
combined_tables = combined_table_dict["1"]
|
|
|
|
# Should still have 2 tables on page 1
|
|
self.assertEqual(len(combined_tables), 2)
|
|
|
|
# Table A (first table) should be unchanged
|
|
table_a_after = combined_tables[0]
|
|
self.assertEqual(
|
|
table_a_after.table.shape[0], table_a_original_rows
|
|
) # Same row count
|
|
self.assertEqual(table_a_after.header, "Table A (Complete)")
|
|
self.assertEqual(
|
|
table_a_after.table.iloc[0, 0], "A1%"
|
|
) # Original data preserved
|
|
|
|
# Table B (second table) should now include the continuation data
|
|
table_b_after = combined_tables[1]
|
|
expected_combined_rows = table_b_original_rows + continuation_rows # 2 + 2 = 4
|
|
self.assertEqual(table_b_after.table.shape[0], expected_combined_rows)
|
|
self.assertEqual(table_b_after.header, "Table B (Continues)")
|
|
|
|
# Verify Table B has both original and continuation data
|
|
self.assertEqual(table_b_after.table.iloc[0, 0], "B1%") # Original data
|
|
self.assertEqual(table_b_after.table.iloc[1, 0], "B2") # Original data
|
|
self.assertEqual(table_b_after.table.iloc[2, 0], "B3") # Continuation data
|
|
self.assertEqual(table_b_after.table.iloc[3, 0], "B4") # Continuation data
|
|
|
|
def test_combine_continuous_tables_wrong_combination_scenario(self):
|
|
"""Test edge case: what if we accidentally used [0] instead of [-1]?"""
|
|
# This test documents what WOULD happen with the wrong indexing
|
|
# Page 1: Table A + Table B (continues)
|
|
# Page 2: Table B continuation
|
|
# If we used [0], continuation would wrongly combine with Table A
|
|
|
|
text_dict = {
|
|
"1": (
|
|
"-------Table Start--------\nTable A\n[['X', 'Y'], ['X1%', 'Y1']]\n-------Table End--------\n"
|
|
"-------Table Start--------\nTable B\n[['Name', 'Value'], ['B1%', '1']]\n-------Table End--------"
|
|
),
|
|
"2": "-------Table Start--------\n\n[['B2', '2%']]\n-------Table End--------", # Should combine with Table B, not Table A
|
|
}
|
|
|
|
table_dict = get_table_info(text_dict, self.EXHIBIT_HEADER_MARKERS)
|
|
result = combine_continuous_tables(text_dict, table_dict)
|
|
|
|
# Get the result and verify correct combination
|
|
combined_table_dict = get_table_info(
|
|
{"1": result["1"]}, self.EXHIBIT_HEADER_MARKERS
|
|
)
|
|
|
|
# Table A (index 0) should have 1 row still - NOT combined with continuation
|
|
table_a = combined_table_dict["1"][0]
|
|
self.assertEqual(table_a.table.shape[0], 1) # Still just 1 row
|
|
self.assertEqual(table_a.table.iloc[0, 0], "X1%") # Original data only
|
|
|
|
# Table B (index 1, which is -1) should have 2 rows - correctly combined
|
|
table_b = combined_table_dict["1"][1]
|
|
self.assertEqual(table_b.table.shape[0], 2) # 1 original + 1 continuation = 2
|
|
self.assertEqual(table_b.table.iloc[0, 0], "B1%") # Original
|
|
self.assertEqual(table_b.table.iloc[1, 0], "B2") # Continuation
|
|
|
|
def test_table_combine_same_header(self):
|
|
"""Test combining two tables with the same header."""
|
|
table1 = Table(
|
|
page_num="1",
|
|
metadata="metadata",
|
|
header="Table 1",
|
|
table_data=[["Column 1", "Column 2"], ["A", "B"], ["1", "2"]],
|
|
start_index=0,
|
|
end_index=100,
|
|
)
|
|
table2 = Table(
|
|
page_num="1",
|
|
metadata="metadata",
|
|
header="Table 2",
|
|
table_data=[["Column 1", "Column 2"], ["C", "D"], ["3", "4"]],
|
|
start_index=0,
|
|
end_index=100,
|
|
)
|
|
|
|
original_shape = table1.table.shape
|
|
table1.combine(table2)
|
|
self.assertEqual(
|
|
table1.table.shape[0], original_shape[0] + table2.table.shape[0]
|
|
) # Rows should be combined
|
|
self.assertEqual(table1.table.shape, (4, 2)) # 4 rows (2 + 2 rows), 2 columns
|
|
self.assertEqual(table1.table.iloc[0, 0], "A")
|
|
self.assertEqual(
|
|
table1.table.iloc[2, 0], "C"
|
|
) # Check if second table's data is appended correctly
|
|
self.assertEqual(table1.table.iloc[3, 1], "4") # Last row of second table
|
|
|
|
def test_table_combine_different_headers(self):
|
|
"""Test combining two tables with different headers (the second table is assumed to be a continuation in this case).
|
|
We will pass in table_data as STRING for both we use convert_str_to_dataframe and
|
|
check that `is_continuation` is getting instantiated correctly.
|
|
"""
|
|
|
|
table1 = Table(
|
|
page_num="1",
|
|
metadata="metadata",
|
|
header="Table 1",
|
|
table_data="[['Column 1', 'Column 2'], ['A', 'B'], ['1', '2']]",
|
|
start_index=0,
|
|
end_index=100,
|
|
)
|
|
|
|
# Continuation table with empty header
|
|
table2 = Table(
|
|
page_num="1",
|
|
metadata="metadata",
|
|
header="",
|
|
table_data="[['C', 'D'], ['3', '4'], ['E', 'F']]", # Pass as string!
|
|
start_index=0,
|
|
end_index=100,
|
|
is_continuation=True, # Mark as continuation
|
|
)
|
|
|
|
original_shape = table1.table.shape
|
|
table1.combine(table2)
|
|
|
|
print(table1.table)
|
|
self.assertEqual(
|
|
table1.table.shape[0], original_shape[0] + table2.table.shape[0]
|
|
) # Rows should be combined
|
|
self.assertEqual(table1.table.shape, (5, 2)) # 5 rows (2 + 3 rows), 2 columns
|
|
self.assertEqual(table1.table.iloc[0, 0], "A")
|
|
self.assertEqual(
|
|
table1.table.iloc[2, 0], "C"
|
|
) # Check if second table's data is appended correctly
|
|
self.assertEqual(table1.table.iloc[3, 1], "4") # Last row of second table
|
|
self.assertEqual(
|
|
list(table1.table.columns), ["Column 1", "Column 2"]
|
|
) # Should keep the first table's header
|
|
|
|
def test_table_combine_different_columns(self):
|
|
"""Test that combining tables with different column counts handles mismatches gracefully."""
|
|
table1 = Table(
|
|
page_num="1",
|
|
metadata="metadata",
|
|
header="Table 1",
|
|
table_data=[["A", "B"], ["1", "2"]],
|
|
start_index=0,
|
|
end_index=100,
|
|
)
|
|
table2 = Table(
|
|
page_num="1",
|
|
metadata="metadata",
|
|
header="Table 2",
|
|
table_data=[["C", "D", "E"], ["3", "4", "5"]], # 3 columns vs 2 columns
|
|
start_index=0,
|
|
end_index=100,
|
|
)
|
|
|
|
# Should not raise an error anymore - should handle column mismatch
|
|
table1.combine(table2)
|
|
|
|
# Verify the result
|
|
# Should have 2 rows total (1 from each table)
|
|
self.assertEqual(table1.table.shape[0], 2)
|
|
|
|
# Should have 2 columns (from the main table)
|
|
self.assertEqual(table1.table.shape[1], 2)
|
|
|
|
# Column names should match the main table
|
|
self.assertEqual(list(table1.table.columns), ["A", "B"])
|
|
|
|
# First row should be from table1
|
|
self.assertEqual(table1.table.iloc[0, 0], "1")
|
|
self.assertEqual(table1.table.iloc[0, 1], "2")
|
|
|
|
# Second row should be from table2, but truncated to 2 columns
|
|
self.assertEqual(table1.table.iloc[1, 0], "3")
|
|
self.assertEqual(table1.table.iloc[1, 1], "4")
|
|
# The "5" from column "E" should be truncated/lost
|
|
|
|
def test_table_combine_column_padding(self):
|
|
"""Test that combining tables pads shorter tables with empty columns."""
|
|
table1 = Table(
|
|
page_num="1",
|
|
metadata="metadata",
|
|
header="Table 1",
|
|
table_data=[["A", "B", "C"], ["1", "2", "3"]], # 3 columns
|
|
start_index=0,
|
|
end_index=100,
|
|
)
|
|
table2 = Table(
|
|
page_num="1",
|
|
metadata="metadata",
|
|
header="Table 2",
|
|
table_data=[["X"], ["4"]], # Only 1 column
|
|
start_index=0,
|
|
end_index=100,
|
|
)
|
|
|
|
table1.combine(table2)
|
|
|
|
# Should have 2 rows total
|
|
self.assertEqual(table1.table.shape[0], 2)
|
|
|
|
# Should have 3 columns (from the main table)
|
|
self.assertEqual(table1.table.shape[1], 3)
|
|
|
|
# Column names should match the main table
|
|
self.assertEqual(list(table1.table.columns), ["A", "B", "C"])
|
|
|
|
# First row should be unchanged
|
|
self.assertEqual(table1.table.iloc[0, 0], "1")
|
|
self.assertEqual(table1.table.iloc[0, 1], "2")
|
|
self.assertEqual(table1.table.iloc[0, 2], "3")
|
|
|
|
# Second row should be from table2, padded with empty strings
|
|
self.assertEqual(table1.table.iloc[1, 0], "4")
|
|
self.assertEqual(table1.table.iloc[1, 1], "") # Padded
|
|
self.assertEqual(table1.table.iloc[1, 2], "") # Padded
|
|
|
|
def test_table_combine_empty_tables(self):
|
|
"""Test combining with empty tables."""
|
|
table1 = Table(
|
|
page_num="1",
|
|
metadata="meta",
|
|
header="Table 1",
|
|
table_data=[["A", "B"]],
|
|
start_index=0,
|
|
end_index=100,
|
|
)
|
|
empty_table = Table(
|
|
page_num="2",
|
|
metadata="meta",
|
|
header="",
|
|
table_data=[],
|
|
start_index=0,
|
|
end_index=100,
|
|
is_continuation=True,
|
|
)
|
|
|
|
original_shape = table1.table.shape
|
|
table1.combine(empty_table)
|
|
self.assertEqual(table1.table.shape, original_shape) # Should remain unchanged
|
|
|
|
def test_convert_str_to_dataframe_malformed_data(self):
|
|
"""Test handling of malformed table data."""
|
|
malformed_inputs = [
|
|
"[['A', 'B'], ['1']]", # Inconsistent row lengths
|
|
"[['A'], [], ['B']]", # Empty rows
|
|
"[[]]", # Empty nested list
|
|
"['not', 'nested']", # Not nested list
|
|
]
|
|
|
|
for malformed_input in malformed_inputs:
|
|
with self.subTest(input=malformed_input):
|
|
df = convert_str_to_dataframe(malformed_input)
|
|
self.assertIsInstance(df, pd.DataFrame)
|
|
|
|
def test_table_post_table_text_preservation(self):
|
|
"""Test that post-table text is properly preserved through operations."""
|
|
text_dict = {
|
|
"1": "Metadata\n-------Table Start--------\nTable\n[['A', 'B'], ['1', '2%']]\n-------Table End--------\nImportant footnote\nSignature line"
|
|
}
|
|
|
|
table_dict = get_table_info(text_dict, self.EXHIBIT_HEADER_MARKERS)
|
|
table = table_dict["1"][0]
|
|
|
|
# Should capture post-table text
|
|
self.assertIn("Important footnote", table.post_table_text)
|
|
self.assertIn("Signature line", table.post_table_text)
|
|
|
|
def test_split_large_table_preserves_metadata_and_post_text(self):
|
|
"""Test that table splitting preserves all important information."""
|
|
table = Table(
|
|
page_num="31",
|
|
metadata="EXHIBIT F",
|
|
header="Rate Table",
|
|
table_data=[["HCPC", "Rate"]]
|
|
+ [[f"T{i:04d}", f"${i}.00"] for i in range(15)],
|
|
start_index=100,
|
|
end_index=2000,
|
|
post_table_text="*Rates effective 2024",
|
|
)
|
|
|
|
with patch("src.config.TABLE_ROW_LIMIT", 5):
|
|
chunks = split_large_table_with_headers(table)
|
|
|
|
# All chunks should have metadata and header
|
|
for chunk in chunks:
|
|
self.assertIn("EXHIBIT F", chunk)
|
|
self.assertIn("Rate Table", chunk)
|
|
|
|
# Only last chunk should have post-table text
|
|
self.assertIn("*Rates effective 2024", chunks[-1])
|
|
for chunk in chunks[:-1]:
|
|
self.assertNotIn("*Rates effective 2024", chunk)
|
|
|
|
def test_table_split(self):
|
|
"""Test splitting a table."""
|
|
long_data = [["Name", "Value"]] + [
|
|
[f"Row{i}", str(i)] for i in range(1, 6)
|
|
] # This table will have 5 rows
|
|
long_table = Table(
|
|
page_num="1",
|
|
metadata="metadata",
|
|
header="Long Table",
|
|
table_data=long_data,
|
|
start_index=0,
|
|
end_index=100,
|
|
)
|
|
split_tables = long_table.split(row_limit=2)
|
|
|
|
self.assertEqual(
|
|
len(split_tables), 3
|
|
) # Should split into 3 tables (5 rows split by 2 = 3)
|
|
self.assertEqual(
|
|
split_tables[0].table.shape, (2, 2)
|
|
) # First split should have 2 rows
|
|
self.assertEqual(
|
|
split_tables[1].table.shape, (2, 2)
|
|
) # Second split should have 2 rows
|
|
self.assertEqual(
|
|
split_tables[2].table.shape, (1, 2)
|
|
) # Last split should have 1 row
|
|
|
|
# Check that original metadata is preserved
|
|
for split_table in split_tables:
|
|
self.assertEqual(split_table.metadata, "metadata")
|
|
self.assertEqual(split_table.page_num, "1")
|
|
self.assertEqual(split_table.header, "Long Table")
|
|
|
|
def test_table_split_no_split_needed(self):
|
|
"""Test splitting a table that does not need splitting. It should return the original table."""
|
|
table = Table(
|
|
page_num="1",
|
|
metadata="metadata",
|
|
header="Table 1",
|
|
table_data=[["A", "B"], ["1", "2"], ["3", "4"]],
|
|
start_index=0,
|
|
end_index=100,
|
|
)
|
|
split_tables = table.split(row_limit=10) # Row limit larger than table size
|
|
self.assertEqual(len(split_tables), 1) # Should return the original table
|
|
self.assertEqual(
|
|
split_tables[0].table.shape, (2, 2)
|
|
) # Should be the same shape as original
|
|
self.assertEqual(split_tables[0], table) # Should be the same object
|
|
|
|
def test_table_to_list_format(self):
|
|
"""Test conversion back to list format."""
|
|
table_data = [["A", "B"], ["1", "3"], ["2", "4"]]
|
|
table = Table(
|
|
page_num="1",
|
|
metadata="metadata",
|
|
header="Table 1",
|
|
table_data=table_data,
|
|
start_index=0,
|
|
end_index=100,
|
|
)
|
|
list_format = table.to_list_format()
|
|
|
|
self.assertEqual(list_format, table_data)
|
|
self.assertIsInstance(list_format, list)
|
|
self.assertIsInstance(list_format[0], list) # Should be a list of lists
|
|
|
|
def test_table_info_new_format(self):
|
|
"""Test getting table info from new list-of-rows format."""
|
|
table_dict = get_table_info(
|
|
self.sample_text_dict_new, self.EXHIBIT_HEADER_MARKERS
|
|
)
|
|
|
|
self.assertIn("1", table_dict)
|
|
self.assertEqual(len(table_dict["1"]), 1)
|
|
self.assertEqual(table_dict["1"][0].header, "Table 1")
|
|
self.assertEqual(table_dict["1"][0].metadata, "Some metadata")
|
|
self.assertEqual(
|
|
table_dict["1"][0].table.shape, (2, 2)
|
|
) # 2 rows, 2 columns in first table
|
|
self.assertEqual(table_dict["2"][0].header, "") # Second table has no header
|
|
self.assertEqual(
|
|
table_dict["2"][0].table.shape, (3, 2)
|
|
) # 3 rows, 2 columns in second table (empty header means all data rows)
|
|
|
|
def test_table_info_old_format(self):
|
|
"""Test getting table info from old dict format."""
|
|
with self.assertLogs(level="INFO") as log: # Expect a deprecation warning
|
|
table_dict = get_table_info(
|
|
self.sample_text_dict_old, self.EXHIBIT_HEADER_MARKERS
|
|
)
|
|
|
|
self.assertIn("1", table_dict)
|
|
self.assertEqual(len(table_dict["1"]), 1)
|
|
self.assertEqual(table_dict["1"][0].header, "Table 1")
|
|
self.assertEqual(table_dict["1"][0].metadata, "Some metadata")
|
|
self.assertEqual(table_dict["1"][0].table.shape, (2, 2))
|
|
|
|
def test_get_table_info_multiple_tables(self):
|
|
"""Test getting table info from a page with multiple tables."""
|
|
table_dict = get_table_info(self.multi_table_text, self.EXHIBIT_HEADER_MARKERS)
|
|
|
|
self.assertIn("1", table_dict)
|
|
self.assertEqual(len(table_dict["1"]), 1)
|
|
self.assertEqual(table_dict["1"][0].header, "First Table")
|
|
|
|
def test_recreate_page(self):
|
|
"""Test recreating page text from tables."""
|
|
table_dict = get_table_info(
|
|
self.sample_text_dict_new, self.EXHIBIT_HEADER_MARKERS
|
|
)
|
|
page_tables = table_dict["1"]
|
|
recreated_text = recreate_page(page_tables)
|
|
|
|
self.assertIn("Table 1", recreated_text)
|
|
self.assertIn("-------Table Start--------", recreated_text)
|
|
self.assertIn("-------Table End--------", recreated_text)
|
|
self.assertIn("Some metadata", recreated_text)
|
|
|
|
def test_split_tables_by_row_count(self):
|
|
"""Test splitting tables that exceed row limit and adding .0 suffix."""
|
|
# Create a long table that will need splitting
|
|
long_table_data = [["Name", "Value"]] + [
|
|
[f"Row{i}%", str(i)] for i in range(1, 8)
|
|
] # 7 data rows
|
|
text_dict = {
|
|
"5": f"Table metadata\n-------Table Start--------\nLong Table\n{long_table_data}\n-------Table End--------\nPost-table text"
|
|
}
|
|
|
|
result = split_tables_by_rows(
|
|
text_dict, self.EXHIBIT_HEADER_MARKERS, row_limit=3
|
|
)
|
|
|
|
# Should create multiple pages with .0, .1, .2 suffixes
|
|
expected_pages = ["5.0", "5.1", "5.2"]
|
|
for page in expected_pages:
|
|
self.assertIn(page, result)
|
|
|
|
# Original page should not exist
|
|
self.assertNotIn("5", result)
|
|
|
|
# Last page should have post-table text
|
|
self.assertIn("Post-table text", result["5.2"])
|
|
|
|
def test_split_tables_by_row_count_no_split_needed(self):
|
|
"""Test that small tables get .0 suffix but aren't split."""
|
|
text_dict = {
|
|
"3": "Metadata\n-------Table Start--------\nSmall Table\n[['A', 'B'], ['1%', '2']]\n-------Table End--------"
|
|
}
|
|
|
|
result = split_tables_by_rows(
|
|
text_dict, self.EXHIBIT_HEADER_MARKERS, row_limit=10
|
|
)
|
|
|
|
# Should have .0 suffix but no additional splits
|
|
self.assertIn("3.0", result)
|
|
self.assertNotIn("3", result)
|
|
self.assertNotIn("3.1", result)
|
|
|
|
def test_split_tables_by_row_count_no_tables(self):
|
|
"""Test that pages without tables keep original page numbers."""
|
|
text_dict = {
|
|
"1": "Just some regular text without any tables.",
|
|
"2": "More text, no tables here either.",
|
|
}
|
|
|
|
result = split_tables_by_rows(
|
|
text_dict, self.EXHIBIT_HEADER_MARKERS, row_limit=5
|
|
)
|
|
|
|
# Pages without tables should keep original numbers
|
|
self.assertIn("1", result)
|
|
self.assertIn("2", result)
|
|
self.assertNotIn("1.0", result)
|
|
self.assertNotIn("2.0", result)
|
|
|
|
def test_combine_continuous_tables_no_continuation(self):
|
|
"""Test that non-continuation tables are not combined."""
|
|
text_dict = {
|
|
"1": "Metadata 1\n-------Table Start--------\nTable 1\n[['A', 'B'], ['1%', '2']]\n-------Table End--------",
|
|
"2": "Metadata 2\n-------Table Start--------\nTable 2\n[['C', 'D'], ['3%', '4']]\n-------Table End--------",
|
|
}
|
|
|
|
table_dict = get_table_info(text_dict, self.EXHIBIT_HEADER_MARKERS)
|
|
result = combine_continuous_tables(text_dict, table_dict)
|
|
|
|
# Both pages should remain (no continuation tables to combine)
|
|
self.assertIn("1", result)
|
|
self.assertIn("2", result)
|
|
self.assertEqual(len(result), 2)
|
|
|
|
def test_table_combine_post_table_text_accumulation(self):
|
|
"""Test that post-table text is properly accumulated when combining tables."""
|
|
table1 = Table(
|
|
page_num="1",
|
|
metadata="meta",
|
|
header="Main Table",
|
|
table_data=[["A", "B"], ["1", "2"]],
|
|
start_index=0,
|
|
end_index=100,
|
|
post_table_text="First footnote",
|
|
)
|
|
|
|
table2 = Table(
|
|
page_num="2",
|
|
metadata="",
|
|
header="",
|
|
table_data=[["3", "4"]],
|
|
start_index=0,
|
|
end_index=100,
|
|
is_continuation=True,
|
|
post_table_text="Second footnote",
|
|
)
|
|
|
|
table1.combine(table2)
|
|
|
|
# Both footnotes should be preserved
|
|
self.assertIn("First footnote", table1.post_table_text)
|
|
self.assertIn("Second footnote", table1.post_table_text)
|
|
|
|
def test_table_combine_column_mismatch_preserves_post_table_text(self):
|
|
"""Test that post-table text is preserved when combining tables with column mismatches."""
|
|
table1 = Table(
|
|
page_num="1",
|
|
metadata="metadata",
|
|
header="Main Table",
|
|
table_data=[["A", "B"], ["1", "2"], ["x", "y"]], # Header + 2 data rows
|
|
start_index=0,
|
|
end_index=100,
|
|
post_table_text="Main table footnote",
|
|
)
|
|
|
|
table2 = Table(
|
|
page_num="2",
|
|
metadata="",
|
|
header="",
|
|
table_data=[["3", "4", "5"]], # Just 1 data row (3 columns - mismatch!)
|
|
start_index=0,
|
|
end_index=100,
|
|
is_continuation=True,
|
|
post_table_text="Continuation footnote",
|
|
)
|
|
|
|
table1.combine(table2)
|
|
|
|
# Verify both post-table texts are preserved despite column mismatch
|
|
self.assertIn("Main table footnote", table1.post_table_text)
|
|
self.assertIn("Continuation footnote", table1.post_table_text)
|
|
|
|
# Verify the table combination worked correctly
|
|
self.assertEqual(
|
|
table1.table.shape[0], 3
|
|
) # 2 original + 1 continuation = 3 total rows
|
|
self.assertEqual(table1.table.shape[1], 2) # 2 columns (from main table)
|
|
|
|
# Verify data integrity
|
|
self.assertEqual(table1.table.iloc[0, 0], "1") # Original data row 1
|
|
self.assertEqual(table1.table.iloc[1, 0], "x") # Original data row 2
|
|
self.assertEqual(table1.table.iloc[2, 0], "3") # Continuation data (truncated)
|
|
|
|
def test_table_combine_column_mismatch_padding_preserves_post_table_text(self):
|
|
"""Test post-table text preservation when padding columns."""
|
|
table1 = Table(
|
|
page_num="1",
|
|
metadata="metadata",
|
|
header="Main Table",
|
|
table_data=[
|
|
["A", "B", "C"],
|
|
["1", "2", "3"],
|
|
["x", "y", "z"],
|
|
], # Header + 2 data rows
|
|
start_index=0,
|
|
end_index=100,
|
|
post_table_text="Main footnote",
|
|
)
|
|
|
|
table2 = Table(
|
|
page_num="2",
|
|
metadata="",
|
|
header="",
|
|
table_data=[["4"]], # Just 1 data row (1 column - needs padding)
|
|
start_index=0,
|
|
end_index=100,
|
|
is_continuation=True,
|
|
post_table_text="Padded footnote",
|
|
)
|
|
|
|
table1.combine(table2)
|
|
|
|
# Both footnotes should be preserved even with column padding
|
|
self.assertIn("Main footnote", table1.post_table_text)
|
|
self.assertIn("Padded footnote", table1.post_table_text)
|
|
|
|
# Verify padding worked
|
|
self.assertEqual(
|
|
table1.table.shape, (3, 3)
|
|
) # 2 original + 1 continuation = 3 rows, 3 columns
|
|
self.assertEqual(table1.table.iloc[0, 0], "1") # Original data row 1
|
|
self.assertEqual(table1.table.iloc[1, 0], "x") # Original data row 2
|
|
self.assertEqual(table1.table.iloc[2, 0], "4") # Continuation data
|
|
self.assertEqual(table1.table.iloc[2, 1], "") # Padded
|
|
self.assertEqual(table1.table.iloc[2, 2], "") # Padded
|
|
|
|
def test_exhibit_boundary_detection_prevents_combination(self):
|
|
"""Test that exhibit boundaries prevent incorrect table combination."""
|
|
text_dict = {
|
|
"1": "EXHIBIT A-1\nCommercial rates\n-------Table Start--------\n\n[['Service', 'Rate'], ['A', '1%']]\n-------Table End--------",
|
|
"2": "EXHIBIT B-1\nMedicare rates\n-------Table Start--------\n\n[['Service', 'Rate'], ['B', '2%']]\n-------Table End--------",
|
|
}
|
|
|
|
table_dict = get_table_info(text_dict, self.EXHIBIT_HEADER_MARKERS)
|
|
|
|
# Both should be detected as non-continuation tables due to exhibit boundaries
|
|
self.assertFalse(table_dict["1"][0].is_continuation)
|
|
self.assertFalse(table_dict["2"][0].is_continuation)
|
|
|
|
# Combination should not remove any pages
|
|
result = combine_continuous_tables(text_dict, table_dict)
|
|
self.assertIn("1", result)
|
|
self.assertIn("2", result)
|
|
self.assertEqual(len(result), 2)
|
|
|
|
def test_exhibit_boundary_case_insensitive(self):
|
|
"""Test that exhibit detection works with different cases."""
|
|
test_cases = [
|
|
"EXHIBIT A-1",
|
|
"Exhibit A-1",
|
|
"exhibit a-1",
|
|
"ATTACHMENT B",
|
|
"Schedule C",
|
|
"ADDENDUM FOR MEDICARE",
|
|
]
|
|
|
|
for exhibit_text in test_cases:
|
|
with self.subTest(exhibit=exhibit_text):
|
|
text_dict = {
|
|
"1": f"{exhibit_text}\nRate information\n-------Table Start--------\n\n[['A', 'B'], ['1%', '2']]\n-------Table End--------"
|
|
}
|
|
|
|
table_dict = get_table_info(text_dict, self.EXHIBIT_HEADER_MARKERS)
|
|
|
|
# Should not be detected as continuation due to exhibit marker
|
|
self.assertFalse(table_dict["1"][0].is_continuation)
|
|
|
|
def test_exhibit_boundary_detection_beyond_100_chars(self):
|
|
"""Test that exhibit markers beyond 100 characters don't affect continuation detection."""
|
|
# Create a page with exhibit marker after position 100
|
|
long_prefix = "A" * 95 # 95 characters
|
|
text_dict = {
|
|
"1": f"{long_prefix}EXHIBIT A-1\nRate info\n-------Table Start--------\n\n[['A', 'B'], ['1%', '2']]\n-------Table End--------"
|
|
}
|
|
|
|
table_dict = get_table_info(text_dict, self.EXHIBIT_HEADER_MARKERS)
|
|
|
|
# Should be treated as continuation since exhibit marker is beyond position 100
|
|
self.assertTrue(table_dict["1"][0].is_continuation)
|
|
|
|
def test_table_init_with_empty_list(self):
|
|
"""Test Table initialization with empty list."""
|
|
table = Table(
|
|
page_num="1",
|
|
metadata="meta",
|
|
header="Empty Table",
|
|
table_data=[],
|
|
start_index=0,
|
|
end_index=100,
|
|
)
|
|
|
|
self.assertTrue(table.table.empty)
|
|
self.assertEqual(table.table.shape, (0, 0))
|
|
|
|
def test_table_init_with_only_headers(self):
|
|
"""Test Table initialization with only header row (no data)."""
|
|
table = Table(
|
|
page_num="1",
|
|
metadata="meta",
|
|
header="Header Only",
|
|
table_data=[["Column A", "Column B"]],
|
|
start_index=0,
|
|
end_index=100,
|
|
)
|
|
|
|
self.assertEqual(table.table.shape, (0, 2)) # 0 rows, 2 columns
|
|
self.assertEqual(list(table.table.columns), ["Column A", "Column B"])
|
|
|
|
def test_continuation_table_with_generic_column_names_combination(self):
|
|
"""Test that continuation tables with generic column names combine correctly."""
|
|
main_table = Table(
|
|
page_num="1",
|
|
metadata="meta",
|
|
header="Main Table",
|
|
table_data=[["Service", "Rate"], ["A", "1%"], ["B", "2%"]],
|
|
start_index=0,
|
|
end_index=100,
|
|
)
|
|
|
|
# Continuation table with generic column names
|
|
continuation_table = Table(
|
|
page_num="2",
|
|
metadata="",
|
|
header="",
|
|
table_data=[["C", "3%"], ["D", "4%"]],
|
|
start_index=0,
|
|
end_index=100,
|
|
is_continuation=True,
|
|
)
|
|
|
|
main_table.combine(continuation_table)
|
|
|
|
# Should have 4 data rows total
|
|
self.assertEqual(main_table.table.shape, (4, 2))
|
|
|
|
# Column names should remain from main table
|
|
self.assertEqual(list(main_table.table.columns), ["Service", "Rate"])
|
|
|
|
# Data should be properly combined
|
|
self.assertEqual(main_table.table.iloc[0, 0], "A") # Original
|
|
self.assertEqual(main_table.table.iloc[2, 0], "C") # Continuation
|
|
self.assertEqual(main_table.table.iloc[3, 1], "4%") # Last continuation row
|
|
|
|
def test_recreate_page_with_multiple_tables_and_post_text(self):
|
|
"""Test page recreation with multiple tables and post-table text."""
|
|
tables = [
|
|
Table(
|
|
page_num="1",
|
|
metadata="Page header",
|
|
header="First Table",
|
|
table_data=[["A", "B"], ["1%", "2"]],
|
|
start_index=0,
|
|
end_index=100,
|
|
),
|
|
Table(
|
|
page_num="1",
|
|
metadata="",
|
|
header="Second Table",
|
|
table_data=[["X", "Y"], ["3%", "4"]],
|
|
start_index=200,
|
|
end_index=300,
|
|
post_table_text="Important footnote\nSignature required",
|
|
),
|
|
]
|
|
|
|
recreated = recreate_page(tables)
|
|
|
|
# Should contain both tables
|
|
self.assertIn("First Table", recreated)
|
|
self.assertIn("Second Table", recreated)
|
|
|
|
# Should contain post-table text from last table
|
|
self.assertIn("Important footnote", recreated)
|
|
self.assertIn("Signature required", recreated)
|
|
|
|
# Should have proper structure
|
|
self.assertIn(START_MARKER, recreated)
|
|
self.assertIn(END_MARKER, recreated)
|
|
|
|
def test_split_tables_by_rows_preserves_metadata_and_post_text(self):
|
|
"""Test that table splitting preserves all text components."""
|
|
table_data = [["Code", "Rate"]] + [[f"T{i:04d}", f"${i}.00"] for i in range(12)]
|
|
|
|
text_dict = {
|
|
"5": f"EXHIBIT F\nRate Schedule\n-------Table Start--------\nLarge Table\n{table_data}\n-------Table End--------\nEffective 2024\nContact: admin@health.net"
|
|
}
|
|
|
|
result = split_tables_by_rows(
|
|
text_dict, self.EXHIBIT_HEADER_MARKERS, row_limit=5
|
|
)
|
|
|
|
# Should create multiple pages
|
|
expected_pages = ["5.0", "5.1", "5.2"]
|
|
for page in expected_pages:
|
|
self.assertIn(page, result)
|
|
|
|
# All pages should have metadata and header
|
|
for page in expected_pages:
|
|
self.assertIn("EXHIBIT F", result[page])
|
|
self.assertIn("Large Table", result[page])
|
|
|
|
# Only last page should have post-table text
|
|
self.assertIn("Effective 2024", result["5.2"])
|
|
self.assertIn("Contact: admin@health.net", result["5.2"])
|
|
|
|
# Earlier pages should not have post-table text
|
|
self.assertNotIn("Effective 2024", result["5.0"])
|
|
self.assertNotIn("Effective 2024", result["5.1"])
|
|
|
|
def test_convert_str_to_dataframe_edge_cases(self):
|
|
"""Test edge cases in string to DataFrame conversion."""
|
|
edge_cases = [
|
|
("[]", True), # Empty list
|
|
("[['A']]", False), # Single cell
|
|
("[['A', 'B'], []]", False), # Empty data row
|
|
("[[], ['1', '2']]", False), # Empty header row
|
|
("[['']]", False), # Single empty string
|
|
]
|
|
|
|
for test_input, is_continuation in edge_cases:
|
|
with self.subTest(input=test_input, continuation=is_continuation):
|
|
df = convert_str_to_dataframe(
|
|
test_input, is_continuation=is_continuation
|
|
)
|
|
self.assertIsInstance(df, pd.DataFrame)
|
|
# Should not raise exceptions
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|