From 896918570e84e312a7d7835413fe9db146261d0b Mon Sep 17 00:00:00 2001 From: Matthias Bisping Date: Thu, 10 Mar 2022 16:09:40 +0100 Subject: [PATCH] added response strategies to queue visitor --- config.yaml | 2 +- pyinfra/test/unit_tests/conftest.py | 9 +++++ pyinfra/test/unit_tests/queue_visitor_test.py | 13 +++---- pyinfra/visitor.py | 35 +++++++++++++++++-- 4 files changed, 50 insertions(+), 9 deletions(-) diff --git a/config.yaml b/config.yaml index b03a504..89ac8a1 100755 --- a/config.yaml +++ b/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 diff --git a/pyinfra/test/unit_tests/conftest.py b/pyinfra/test/unit_tests/conftest.py index 328a76a..f39785a 100644 --- a/pyinfra/test/unit_tests/conftest.py +++ b/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() diff --git a/pyinfra/test/unit_tests/queue_visitor_test.py b/pyinfra/test/unit_tests/queue_visitor_test.py index 8f1aba3..f667d65 100644 --- a/pyinfra/test/unit_tests/queue_visitor_test.py +++ b/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" diff --git a/pyinfra/visitor.py b/pyinfra/visitor.py index a011d09..ac7457e 100644 --- a/pyinfra/visitor.py +++ b/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)