18 lines
477 B
Python
18 lines
477 B
Python
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]
|