From 1b1da50faf3c5a0a9c094c5bd37fc5769806d63d Mon Sep 17 00:00:00 2001
From: Julius Unverfehrt <Julius.Unverfehrt@iqser.com>
Date: Mon, 21 Mar 2022 15:22:38 +0100
Subject: [PATCH] Pull request #24: Response dict & Prometheus forwarding

Merge in RR/pyinfra from response-dict to master

Squashed commit of the following:

commit aed09653aafe7b9558c542ea0c36a5f072c6ab9e
Author: Julius Unverfehrt <julius.unverfehrt@iqser.com>
Date:   Mon Mar 21 15:20:24 2022 +0100

    BLACK

commit d1b76673a22373c32815466c37dcb006fa0d6fb4
Author: Julius Unverfehrt <julius.unverfehrt@iqser.com>
Date:   Mon Mar 21 15:19:55 2022 +0100

    change default response dict

commit 82bd48a7545b9b244bdcdebf7ce976947eecf580
Author: Julius Unverfehrt <julius.unverfehrt@iqser.com>
Date:   Mon Mar 21 15:19:15 2022 +0100

    fix handling of missing prometheus endpoint
---
 pyinfra/flask.py       | 16 ++++++++++------
 scripts/mock_client.py |  3 ++-
 src/serve.py           |  2 ++
 3 files changed, 14 insertions(+), 7 deletions(-)

diff --git a/pyinfra/flask.py b/pyinfra/flask.py
index 8196c07..fcddb97 100644
--- a/pyinfra/flask.py
+++ b/pyinfra/flask.py
@@ -50,11 +50,15 @@ def set_up_probing_webserver():
         try:
             resp = requests.get(f"{CONFIG.rabbitmq.callback.analysis_endpoint}/prometheus")
             resp.raise_for_status()
-            return resp.text
-        except Exception as err:
-            if not informed_about_missing_prometheus_endpoint:
-                logger.warning(f"Got no metrics from analysis prometheus endpoint: {err}")
-                informed_about_missing_prometheus_endpoint = True
-            return resp
+        except ConnectionError:
+            return ""
+        except requests.exceptions.HTTPError as err:
+            if resp.status_code == 404:
+                if not informed_about_missing_prometheus_endpoint:
+                    logger.warning(f"Got no metrics from analysis prometheus endpoint: {err}")
+                    informed_about_missing_prometheus_endpoint = True
+            else:
+                raise err
+        return resp.text
 
     return app
diff --git a/scripts/mock_client.py b/scripts/mock_client.py
index d13b0c9..89b2b90 100644
--- a/scripts/mock_client.py
+++ b/scripts/mock_client.py
@@ -48,7 +48,8 @@ def build_message_bodies(analyse_container_type):
             message_dict.update({"targetFileExtension": "ORIGIN.pdf.gz", "responseFileExtension": "IMAGE_INFO.json.gz"})
         if analyse_container_type == "ner":
             message_dict.update(
-                {"targetFileExtension": "TEXT.json.gz", "responseFileExtension": "NER_ENTITIES.json.gz"})
+                {"targetFileExtension": "TEXT.json.gz", "responseFileExtension": "NER_ENTITIES.json.gz"}
+            )
         return message_dict
 
     storage = get_s3_storage()
diff --git a/src/serve.py b/src/serve.py
index 8b76724..1948abf 100644
--- a/src/serve.py
+++ b/src/serve.py
@@ -31,6 +31,8 @@ def make_callback(analysis_endpoint):
         operations = message.get("operations", ["/"])
         results = map(perform_operation, operations)
         result = dict(zip(operations, results))
+        if list(result.keys()) == ["/"]:
+            result = list(result.values())[0]
         return result
 
     return callback