diff --git a/pyinfra/test/unit_tests/conftest.py b/pyinfra/test/unit_tests/conftest.py index be06bf2..915df72 100644 --- a/pyinfra/test/unit_tests/conftest.py +++ b/pyinfra/test/unit_tests/conftest.py @@ -94,13 +94,15 @@ def queue_manager(queue_manager_name): @pytest.fixture(scope="session") def callback(): - return lambda x: x * 2 + def inner(request): + return request["data"].decode() * 2 + return inner @pytest.fixture def analysis_callback(callback): - def inner(data: bytes): - return callback(data.decode()) + def inner(request): + return callback(request) return inner diff --git a/pyinfra/test/unit_tests/consumer_test.py b/pyinfra/test/unit_tests/consumer_test.py index f30996b..7f9705c 100644 --- a/pyinfra/test/unit_tests/consumer_test.py +++ b/pyinfra/test/unit_tests/consumer_test.py @@ -1,4 +1,5 @@ import gzip +import json import logging from operator import itemgetter @@ -49,6 +50,13 @@ class TestConsumer: def produce_items(): return map(str, range(3)) + def mock_visitor(callback): + def inner(data): + return callback({"data": data.encode()}) + return inner + + callback = mock_visitor(callback) + queue_manager.clear() for item in produce_items(): diff --git a/pyinfra/visitor.py b/pyinfra/visitor.py index 1d98eee..9e0a30f 100644 --- a/pyinfra/visitor.py +++ b/pyinfra/visitor.py @@ -75,11 +75,11 @@ class QueueVisitor: logging.warning(f"Loading data from storage failed for {object_descriptor}.") raise DataLoadingFailure() from err - def process_data(self, data): - return self.callback(data) + def process_data(self, data, body): + return self.callback({**body, "data": data}) def load_and_process(self, body): - data = self.process_data(self.load_data(body)) + data = self.process_data(self.load_data(body), body) result_body = {**body, "data": data} return result_body