added response strategies to queue visitor

config.yaml

@@ -3,7 +3,7 @@ service:
   response:
     type: $RESPONSE_TYPE|"file"  # Whether the analysis response is stored as file on storage or sent as stream
     extension: $RESPONSE_FILE_EXTENSION|"IMAGE_INFO.json.gz"  # {.IMAGE_INFO.json.gz | .NER_ENTITIES.json.gz}
-    key: $RESPONSE_KEY|"imageMetadata" # the key of the result {result, imageMetadata}
+#    key: $RESPONSE_KEY|"imageMetadata" # the key of the result {result, imageMetadata}
 
 probing_webserver:
   host: $PROBING_WEBSERVER_HOST|"0.0.0.0"  # Probe webserver address

pyinfra/test/unit_tests/conftest.py

@@ -17,6 +17,7 @@ from pyinfra.test.config import CONFIG
 from pyinfra.test.queue_manager_mock import QueueManagerMock
 from pyinfra.test.storage.adapter_mock import StorageAdapterMock
 from pyinfra.test.storage.client_mock import StorageClientMock
+from pyinfra.visitor import StorageStrategy, ForwardingStrategy
 
 logger = logging.getLogger(__name__)
 logger.setLevel(logging.DEBUG)
@@ -98,3 +99,11 @@ def analysis_callback(callback):
     def inner(data: bytes):
         return callback(data.decode())
     return inner
+
+
+@pytest.fixture
+def response_strategy(response_strategy_name):
+    if response_strategy_name == "storage":
+        return StorageStrategy()
+    if response_strategy_name == "forwarding":
+        return ForwardingStrategy()

pyinfra/test/unit_tests/queue_visitor_test.py

@@ -4,8 +4,8 @@ from pyinfra.visitor import QueueVisitor, get_object_descriptor
 
 
 @pytest.fixture()
-def visitor(storage, analysis_callback):
-    return QueueVisitor(storage, analysis_callback)
+def visitor(storage, analysis_callback, response_strategy):
+    return QueueVisitor(storage, analysis_callback, response_strategy)
 
 
 @pytest.fixture()
@@ -14,7 +14,8 @@ def body():
 
 
 @pytest.mark.parametrize("client_name", ["mock", "azure", "s3"], scope="session")
-class TestStorage:
+@pytest.mark.parametrize("response_strategy_name", ["forwarding", "storage"], scope="session")
+class TestVisitor:
 
     def test_given_a_input_queue_message_callback_pulls_the_data_from_storage(self, visitor, body, storage, bucket_name):
         storage.clear_bucket(bucket_name)
@@ -22,8 +23,8 @@ class TestStorage:
         data_received = visitor.load_data(body)
         assert b"content" == data_received
 
-    def test_visitor_call_method_pulls_data_and_executes_callback(self, visitor, body, storage, bucket_name):
+    def test_visitor_pulls_and_processes_data(self, visitor, body, storage, bucket_name):
         storage.clear_bucket(bucket_name)
         storage.put_object(**get_object_descriptor(body), data="2".encode())
-        response_body = visitor(body)
-        assert response_body == "22"
+        response_body = visitor.load_and_process(body)
+        assert response_body["data"] == "22"

pyinfra/visitor.py

@@ -1,3 +1,5 @@
+import abc
+import json
 from operator import itemgetter
 from typing import Callable
 
@@ -19,11 +21,34 @@ def get_object_descriptor(body):
     }
 
 
+class ResponseStrategy(abc.ABC):
+    @abc.abstractmethod
+    def handle_response(self, body):
+        pass
+
+    def __call__(self, body):
+        return self.handle_response(body)
+
+
+class StorageStrategy(ResponseStrategy):
+    def handle_response(self, body):
+        self.storage.put_object(**get_object_descriptor(body), data=json.dumps(body.pop("data")))
+        return body
+
+
+class ForwardingStrategy(ResponseStrategy):
+    def handle_response(self, body):
+        return body
+
+
 class QueueVisitor:
 
-    def __init__(self, storage: Storage, callback: Callable):
+    def __init__(self, storage: Storage, callback: Callable, response_strategy):
         self.storage = storage
         self.callback = callback
+        self.response_strategy = response_strategy
+        # TODO: a little dirty
+        self.response_strategy.storage = self.storage
 
     def load_data(self, body):
         return self.storage.get_object(**get_object_descriptor(body))
@@ -31,5 +56,11 @@ class QueueVisitor:
     def process_data(self, data):
         return self.callback(data)
 
+    def load_and_process(self, body):
+        data = self.process_data(self.load_data(body))
+        result_body = {**get_object_descriptor(body), "data": data}
+        return result_body
+
     def __call__(self, body):
-        return self.process_data(self.load_data(body))
+        result_body = self.load_and_process(body)
+        return self.response_strategy(result_body)