Katon Minhas
86db2ea0b7
Bugfix/table error fix * handle multipage table data correctly * same metadata for all tables on the page * Multiple tables on page, non-continuous, split * Add back align_and_format_tables * update unit tests * Merged main into bugfix/table-error-fix * improved align and format table prompt * Successful E2E test on San Joaquin - without running align and format table prompt * San Joaquin e2e test passed * Merged main into bugfix/table-error-fix * Merge branch 'main' into bugfix/table-error-fix * Add same number of columns condition * Merged main into bugfix/table-error-fix * Clean investment_values * Remove import * docstrings * Remove test doc * Add back original ALIGN_AND_FORMAT_TABLES prompt for client work * Remove import * Add config reference * Merged main into bugfix/table-error-fix * Move table_utils to investment.table_funcs * restore table_utils.py * Update tests * update preprocess Approved-by: Alex Galarce
50 lines
2.4 KiB
Python
50 lines
2.4 KiB
Python
import unittest
|
|
from src.investment.table_funcs import Table, get_table_info, recreate_page, remove_table_from_page, combine_continuous_tables, split_long_tables
|
|
|
|
class TestTableFunctions(unittest.TestCase):
|
|
def setUp(self):
|
|
self.sample_text_dict = {
|
|
"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--------"
|
|
}
|
|
|
|
self.sample_table_dict = get_table_info(self.sample_text_dict)
|
|
|
|
def test_get_table_info(self):
|
|
self.assertIn("1", self.sample_table_dict)
|
|
self.assertEqual(len(self.sample_table_dict["1"]), 1)
|
|
self.assertEqual(self.sample_table_dict["1"][0].header, "Table 1")
|
|
|
|
def test_recreate_page(self):
|
|
page_tables = self.sample_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)
|
|
|
|
def test_remove_table_from_page(self):
|
|
table = self.sample_table_dict["1"][0]
|
|
updated_text = remove_table_from_page(self.sample_text_dict["1"], table)
|
|
self.assertNotIn("Table 1", updated_text)
|
|
self.assertNotIn("-------Table Start--------", updated_text)
|
|
|
|
def test_combine_continuous_tables(self):
|
|
combined_text_dict = combine_continuous_tables(self.sample_text_dict, self.sample_table_dict)
|
|
self.assertIn("1", combined_text_dict)
|
|
self.assertNotIn("2", combined_text_dict)
|
|
self.assertIn("5", combined_text_dict["1"]) # Ensuring second table is merged
|
|
|
|
def test_split_long_tables(self):
|
|
long_table_dict = {
|
|
"1": "Some metadata\n-------Table Start--------\nTable 1\n{'A': ['1', '2', '3', '4', '5'], 'B': ['6', '7', '8', '9', '10']}\n-------Table End--------"
|
|
}
|
|
table_dict = get_table_info(long_table_dict)
|
|
split_text_dict = split_long_tables(long_table_dict, table_dict, row_limit=2)
|
|
self.assertGreater(len(split_text_dict), 1)
|
|
for key, text in split_text_dict.items():
|
|
self.assertIn("-------Table Start--------", text)
|
|
self.assertIn("-------Table End--------", text)
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|