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