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()
|