from itertools import chain

import pytest

from pyinfra.server.bufferizer.lazy_bufferizer import StreamBuffer
from pyinfra.server.dispatcher.dispatcher import Nothing
from pyinfra.utils.func import lift


@pytest.mark.parametrize("buffer_size", [0, 1, 3, 10, 12])
def test_stream_buffer(buffer_size):
    def func(x):
        return x ** 2

    func = StreamBuffer(lift(func))

    assert list(chain(*map(func, [*range(10), Nothing]))) == [0, 1, 4, 9, 16, 25, 36, 49, 64, 81]