from funcy import compose, lmapcat

from pyinfra.utils.buffer import bufferize


def test_buffer():
    def buffer_mean(xs):
        return [sum(xs) / len(xs)] if xs else []

    buffer_mean = bufferize(compose(buffer_mean, list), buffer_size=3)
    ys = lmapcat(buffer_mean, range(20))
    assert list(ys) == [1.0, 4.0, 7.0, 10.0, 13.0, 16.0]

    def reverse_buffer(xs):
        return reversed(xs)

    reverse_buffer = bufferize(compose(reverse_buffer, list), buffer_size=3)
    ys = lmapcat(reverse_buffer, range(10))
    assert ys == [2, 1, 0, 5, 4, 3, 8, 7, 6]

    def buffer_sum(xs):
        return [sum(xs)]

    buffer_sum = bufferize(buffer_sum, buffer_size=3)
    ys = lmapcat(buffer_sum, range(10))
    assert ys == [0, 0, 3, 0, 0, 12, 0, 0, 21, 0]