From 963132afb17eb8eb2f877f3ef6d4662d92bce7cc Mon Sep 17 00:00:00 2001
From: Matthias Bisping <matthias.bisping@iqser.com>
Date: Thu, 10 Mar 2022 16:41:53 +0100
Subject: [PATCH] added test for side-effect of storage response strategy

---
 config.yaml                                   |  3 ++-
 pyinfra/test/unit_tests/queue_visitor_test.py | 15 ++++++++++++++-
 pyinfra/visitor.py                            | 18 ++++++++++++++++--
 3 files changed, 32 insertions(+), 4 deletions(-)

diff --git a/config.yaml b/config.yaml
index 89ac8a1..b17a693 100755
--- a/config.yaml
+++ b/config.yaml
@@ -2,7 +2,7 @@ service:
   logging_level: $LOGGING_LEVEL_ROOT|INFO  # Logging level for service logger
   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}
+#    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}
 
 probing_webserver:
@@ -39,6 +39,7 @@ storage:
   bucket: $STORAGE_BUCKET|"pyinfra-test-bucket"  # The bucket / container to pull files specified in queue requests from
   # TODO: Caller should specify exact file name, including extension!
   target_file_extension: $TARGET_FILE_EXTENSION|"ORIGIN.pdf.gz"  # {.TEXT.json.gz | .ORIGIN.pdf.gz} Defines type of file to pull from storage
+  response_file_extension: $RESPONSE_FILE_EXTENSION|"IMAGE_INFO.json.gz"
 
   s3:
     endpoint: $STORAGE_ENDPOINT|"http://127.0.0.1:9000"
diff --git a/pyinfra/test/unit_tests/queue_visitor_test.py b/pyinfra/test/unit_tests/queue_visitor_test.py
index f667d65..97f20f7 100644
--- a/pyinfra/test/unit_tests/queue_visitor_test.py
+++ b/pyinfra/test/unit_tests/queue_visitor_test.py
@@ -1,6 +1,8 @@
+import json
+
 import pytest
 
-from pyinfra.visitor import QueueVisitor, get_object_descriptor
+from pyinfra.visitor import QueueVisitor, get_object_descriptor, StorageStrategy, get_response_object_descriptor
 
 
 @pytest.fixture()
@@ -28,3 +30,14 @@ class TestVisitor:
         storage.put_object(**get_object_descriptor(body), data="2".encode())
         response_body = visitor.load_and_process(body)
         assert response_body["data"] == "22"
+
+    def test_visitor_pulls_and_processes_data(self, visitor, body, storage, bucket_name):
+
+        visitor.response_strategy = StorageStrategy()
+        visitor.response_strategy.storage = storage
+
+        storage.clear_bucket(bucket_name)
+        storage.put_object(**get_object_descriptor(body), data="2".encode())
+        response_body = visitor(body)
+        assert "data" not in response_body
+        assert json.loads(storage.get_object(**get_response_object_descriptor(body))) == "22"
diff --git a/pyinfra/visitor.py b/pyinfra/visitor.py
index ac7457e..af298c2 100644
--- a/pyinfra/visitor.py
+++ b/pyinfra/visitor.py
@@ -13,6 +13,12 @@ def get_object_name(body):
     return object_name
 
 
+def get_response_object_name(body):
+    dossier_id, file_id = itemgetter("dossierId", "fileId")(body)
+    object_name = f"{dossier_id}/{file_id}.{CONFIG.response_file_extension}"
+    return object_name
+
+
 def get_object_descriptor(body):
     return {
         "bucket_name": CONFIG.storage.bucket,
@@ -21,6 +27,14 @@ def get_object_descriptor(body):
     }
 
 
+def get_response_object_descriptor(body):
+    return {
+        "bucket_name": CONFIG.storage.bucket,
+        # TODO: send complete filename in request, as discussed!
+        "object_name": get_response_object_name(body)
+    }
+
+
 class ResponseStrategy(abc.ABC):
     @abc.abstractmethod
     def handle_response(self, body):
@@ -32,7 +46,7 @@ class ResponseStrategy(abc.ABC):
 
 class StorageStrategy(ResponseStrategy):
     def handle_response(self, body):
-        self.storage.put_object(**get_object_descriptor(body), data=json.dumps(body.pop("data")))
+        self.storage.put_object(**get_response_object_descriptor(body), data=json.dumps(body.pop("data")).encode())
         return body
 
 
@@ -58,7 +72,7 @@ class QueueVisitor:
 
     def load_and_process(self, body):
         data = self.process_data(self.load_data(body))
-        result_body = {**get_object_descriptor(body), "data": data}
+        result_body = {**body, "data": data}
         return result_body
 
     def __call__(self, body):