From cc0c1bb9990ee4f8a984660cdbd23de90294a715 Mon Sep 17 00:00:00 2001
From: Clarissa Dietrich <Clarissa.Dietrich@iqser.com>
Date: Thu, 3 Mar 2022 10:27:05 +0100
Subject: [PATCH] Pull request #16: refactor save response as file

Merge in RR/pyinfra from response to master

Squashed commit of the following:

commit a3056c0df73091ee766491331f07e3a0377e7887
Author: cdietrich <clarissa.dietrich@iqser.com>
Date:   Thu Mar 3 10:24:54 2022 +0100

    refactor save response as file
---
 .gitignore          |  1 +
 config.yaml         |  5 +++--
 pyinfra/callback.py | 25 +++++++++++++++++--------
 3 files changed, 21 insertions(+), 10 deletions(-)

diff --git a/.gitignore b/.gitignore
index b864b91..c258607 100644
--- a/.gitignore
+++ b/.gitignore
@@ -6,3 +6,4 @@ reports
 pyinfra.egg-info
 bamboo-specs/target
 .pytest_cache
+/.coverage
diff --git a/config.yaml b/config.yaml
index 6cf2147..ed2c651 100755
--- a/config.yaml
+++ b/config.yaml
@@ -1,8 +1,9 @@
 service:
   logging_level: $LOGGING_LEVEL_ROOT|DEBUG  # Logging level for service logger
   response:
-    type: $RESPONSE_TYPE|"stream"  # Whether the analysis response is stored as file on storage or sent as stream
-    extension: $RESPONSE_FILE_EXTENSION|".NER_ENTITIES.json.gz"  # {.IMAGE_INFO.json.gz | .NER_ENTITIES.json.gz}
+    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}
 
 probing_webserver:
   host: $PROBING_WEBSERVER_HOST|"0.0.0.0"  # Probe webserver address
diff --git a/pyinfra/callback.py b/pyinfra/callback.py
index 6f71d29..c130384 100644
--- a/pyinfra/callback.py
+++ b/pyinfra/callback.py
@@ -59,21 +59,30 @@ def make_callback_for_output_queue(json_wrapped_body_processor, output_queue_nam
     declare_queue(channel, output_queue_name)
 
     def callback(channel, method, _, body):
+        """
+        response is dossier_id, file_id and analysis result, if CONFIG.service.response.type == "file" the response only
+         contains file_id and dossier_id and the analysis result will be written in a json file
+        Args:
+            channel:
+            method:
+            _:
+            body:
+
+        Returns:
+
+        """
 
         dossier_id, file_id, result = json_wrapped_body_processor(body)
 
-        # TODO Unify analysis Repsonse for image-prediction and ner-prediction
-        if CONFIG.service.response.type == "stream":
-            result = {"dossierId": dossier_id, "fileId": file_id, "imageMetadata": result}
-            result = json.dumps(result)
-        elif CONFIG.service.response.type == "file":
-            result = json.dumps(result)
+        result_key = CONFIG.service.response.key if CONFIG.service.response.key else "result"
+        result = json.dumps({"dossierId": dossier_id, "fileId": file_id, result_key: result})
+
+        if CONFIG.service.response.type == "file":
             upload_compressed_response(
                 get_storage(CONFIG.storage.backend), CONFIG.storage.bucket, dossier_id, file_id, result
             )
             result = json.dumps({"dossierId": dossier_id, "fileId": file_id})
-        else:
-            result = "WRONG WRONG WORNG DONG"
+
         channel.basic_publish(exchange="", routing_key=output_queue_name, body=result)
         channel.basic_ack(delivery_tag=method.delivery_tag)