changed callback signature to expect a dict with a key 'data'

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
 

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():

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