from src.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()