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