refactorig of webserver startup logic

config.yaml

@@ -24,11 +24,12 @@ minio:
   bucket: $STORAGE_BUCKET_NAME|redaction  # MinIO bucket
 
 service:
-  logging_level: $LOGGING_LEVEL_ROOT|DEBUG  # Logging level for log file messages
+  logging_level: $LOGGING_LEVEL_ROOT|DEBUG  # Logging level for service logger
   name: $SERVICE_NAME|pyinfra-service-v1  # Name of the service in the kubernetes cluster
   storage_backend: $STORAGE_BACKEND|s3  # The storage to pull files to be processed from
   analysis_endpoint: $ANALYSIS_ENDPOINT|"http://127.0.0.1:5000"
 
-webserver:
-  host: $PROBE_SERVER_HOST|"0.0.0.0"  # Probe webserver address
-  port: $PROBE_SERVER_PORT|8080  # Probe webserver port
+probing_webserver:
+  host: $PROBING_WEBSERVER_HOST|"0.0.0.0"  # Probe webserver address
+  port: $PROBING_WEBSERVER_PORT|8080  # Probe webserver port
+  mode: "development"

pyinfra/flask.py

@@ -0,0 +1,41 @@
+from flask import Flask, jsonify
+from waitress import serve
+
+from pyinfra.config import CONFIG
+
+
+def run_probing_webserver(app, host=None, port=None, mode=None):
+
+    if not host:
+        host = CONFIG.probing_webserver.host
+
+    if not port:
+        port = CONFIG.probing_webserver.port
+
+    if not mode:
+        mode = CONFIG.probing_webserver.port
+
+    if mode == "development":
+        app.run(host=host, port=port, debug=True)
+
+    elif mode == "production":
+        serve(app, host=host, port=port)
+
+
+def set_up_probing_webserver():
+
+    app = Flask(__name__)
+
+    @app.route("/ready", methods=["GET"])
+    def ready():
+        resp = jsonify("OK")
+        resp.status_code = 200
+        return resp
+
+    @app.route("/health", methods=["GET"])
+    def healthy():
+        resp = jsonify("OK")
+        resp.status_code = 200
+        return resp
+
+    return app

src/serve.py

@@ -2,8 +2,6 @@ import logging
 from multiprocessing import Process
 
 import pika
-from flask import Flask, jsonify
-from waitress import serve
 
 from pyinfra.callback import (
     make_retry_callback_for_output_queue,
@@ -14,31 +12,12 @@ from pyinfra.config import CONFIG
 from pyinfra.consume import consume, ConsumerError
 from pyinfra.core import make_payload_processor, make_storage_data_loader, make_analyzer
 from pyinfra.exceptions import UnknownStorageBackend
+from pyinfra.flask import run_probing_webserver, set_up_probing_webserver
 from pyinfra.storage.azure_blob_storage import AzureBlobStorageHandle
 from pyinfra.storage.minio import MinioHandle
 
 
 # TODO: implement meaningful checks
-def start_integrity_checks_webserver(mode="debug"):
-
-    app = Flask(__name__)
-
-    @app.route("/ready", methods=["GET"])
-    def ready():
-        resp = jsonify("OK")
-        resp.status_code = 200
-        return resp
-
-    @app.route("/health", methods=["GET"])
-    def healthy():
-        resp = jsonify("OK")
-        resp.status_code = 200
-        return resp
-
-    if mode == "debug":
-        app.run(host=CONFIG.webserver.host, port=CONFIG.webserver.port, debug=True)
-    elif mode == "production":
-        serve(app, host=CONFIG.webserver.host, port=CONFIG.webserver.port)
 
 
 def get_storage():
@@ -87,8 +66,7 @@ def make_callback():
 
 
 def main():
-
-    webserver = Process(target=start_integrity_checks_webserver, args=("production",))
+    webserver = Process(target=run_probing_webserver, args=(set_up_probing_webserver(),))
     logging.info("Starting webserver...")
     webserver.start()