diff --git a/pyinfra/callback.py b/pyinfra/callback.py index 2392924..754258c 100644 --- a/pyinfra/callback.py +++ b/pyinfra/callback.py @@ -35,7 +35,7 @@ class Callback: """ TODO: Since data and metadata are passed as singletons, there is no buffering and hence no batching happening within the pipeline. However, the queue acknowledgment logic needs to be changed in order to facilitate - passing non-singletons, to only ack a message, once a response is pulled from the output queue of the + passing non-singletons, to only ack a message, once a response is pulled from the output queue of the pipeline. Probably the pipeline return value needs to contains the queue message frame (or so), in order for the queue manager to tell which message to ack. diff --git a/pyinfra/server/buffering/stream.py b/pyinfra/server/buffering/stream.py index e1cd669..1033968 100644 --- a/pyinfra/server/buffering/stream.py +++ b/pyinfra/server/buffering/stream.py @@ -1,7 +1,7 @@ from itertools import chain, takewhile from typing import Iterable -from funcy import first, repeatedly, flatten +from funcy import first, repeatedly, mapcat from pyinfra.server.buffering.bufferize import bufferize from pyinfra.server.dispatcher.dispatcher import Nothing, is_not_nothing @@ -16,7 +16,7 @@ class FlatStreamBuffer: def __call__(self, items): items = chain(items, [Nothing]) - yield from flatten(map(self.stream_buffer, items)) + yield from mapcat(self.stream_buffer, items) class StreamBuffer: diff --git a/test/unit_tests/server/pipeline_test.py b/test/unit_tests/server/pipeline_test.py index 0197d95..50d60e6 100644 --- a/test/unit_tests/server/pipeline_test.py +++ b/test/unit_tests/server/pipeline_test.py @@ -1,5 +1,5 @@ import pytest -from funcy import rcompose, compose, project, second, merge, lpluck +from funcy import rcompose, compose, project, merge from pyinfra.server.buffering.stream import FlatStreamBuffer from pyinfra.server.client_pipeline import ClientPipeline @@ -40,9 +40,9 @@ def test_pipeline( for storage_mdt, queue_mdt in zip(metadata, queue_message_metadata) ] - output = compose(llift(unpack), client_pipeline)(input_data_items, metadata) - assert n_items == 0 or len(output) > 0 - assert output == targets + outputs = compose(llift(unpack), client_pipeline)(input_data_items, metadata) + assert n_items == 0 or len(outputs) > 0 + assert outputs == targets @pytest.mark.parametrize("item_type", ["string"])