Katon Minhas
c4e519894b
Refactor/one time io * Delete client work * Streamline imports * Fix list class * Update tests * Clear code funcs test * Fix unit tests * Single read code funcs * Single-read model * Single-load embeddings * Successful E2E test * update unit tests * Update importsg * Reload poetry.lock * remove old exhibit header function * Remove aarete_derived generic function * remove align and format tables * Remove strings to dict * references * Clear code * Move preprocessing_funcs * remove keywords * refactor postprocessing_funcs * Pass unit test - remove qa_qc directory * Black and isort * remove print Approved-by: Alex Galarce
44 lines
1.1 KiB
Python
44 lines
1.1 KiB
Python
from src.investment.preprocessing_funcs import chunk_by_exhibit
|
|
|
|
|
|
def test1():
|
|
"""Test non-empty text_dict with empty exhibit_pages."""
|
|
text_dict = {"1": "This", "2": "is", "3": "an", "4": "exhibit", "5": "test"}
|
|
exhibit_pages = []
|
|
assert chunk_by_exhibit(text_dict, exhibit_pages) == {
|
|
"1": "1",
|
|
"2": "2",
|
|
"3": "3",
|
|
"4": "4",
|
|
"5": "5",
|
|
}
|
|
|
|
|
|
def test2():
|
|
"""Test when exhibit_pages contains some keys that are also in text_dict."""
|
|
text_dict = {"1": "Welcome", "2": "to", "3": "the", "4": "exhibit"}
|
|
exhibit_pages = ["1", "2", "4"]
|
|
expected_output = {"1": "1", "2": "2", "3": "2", "4": "4"}
|
|
assert chunk_by_exhibit(text_dict, exhibit_pages) == expected_output
|
|
|
|
|
|
def test3():
|
|
"""Test multiple exhibit pages with transitions between them."""
|
|
text_dict = {
|
|
"1": "Page 1",
|
|
"2": "Page 2",
|
|
"3": "Page 3",
|
|
"4": "Page 4",
|
|
"5": "Page 5",
|
|
}
|
|
exhibit_pages = ["2", "4"]
|
|
expected_output = {"1": "0", "2": "2", "3": "2", "4": "4", "5": "4"}
|
|
assert chunk_by_exhibit(text_dict, exhibit_pages) == expected_output
|
|
|
|
|
|
test1()
|
|
|
|
test2()
|
|
|
|
test3()
|