from functools import partial

from funcy import chunks, first, compose


def test_repeated_first_chunk_consumption():
    def f(chunk):
        return sum(chunk)

    def g():
        return f(first(chunks(3, items)))

    items = iter(range(10))

    assert g() == 3
    assert g() == 12
    assert g() == 21
    assert g() == 9


def test_repeated_first_chunk_consumption_passing():
    def f(chunk):
        return sum(chunk)

    g = compose(f, first, partial(chunks, 3))

    items = iter(range(10))

    assert g(items) == 3
    assert g(items) == 12
    assert g(items) == 21
    assert g(items) == 9