From 1d09337378b6011c081fa90f530b8ea84d2ffd79 Mon Sep 17 00:00:00 2001 From: Matthias Bisping Date: Sun, 8 May 2022 17:28:55 +0200 Subject: [PATCH] endpoint suffixes passed to stream processor --- pyinfra/server/server.py | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/pyinfra/server/server.py b/pyinfra/server/server.py index 0111e04..b46b569 100644 --- a/pyinfra/server/server.py +++ b/pyinfra/server/server.py @@ -1,5 +1,6 @@ import logging +import flask from flask import Flask, jsonify, request from funcy import compose, flatten @@ -33,12 +34,14 @@ class RestOndDemandProcessor(OnDemandProcessor): class RestStreamProcessor(RestOndDemandProcessor): """Wraps an on-demand-processor. Combine with a webserver that provides the endpoints 'submit' and 'pickup'.""" - def __init__(self, fn): + def __init__(self, fn, submit_suffix="submit", pickup_suffix="pickup"): super(RestStreamProcessor, self).__init__(fn=fn) + self.submit_suffix = submit_suffix + self.pickup_suffix = pickup_suffix def submit(self, request, **kwargs): super(RestStreamProcessor, self).submit(request) - return jsonify(f"{request.base_url.replace('/submit', '')}/pickup") + return jsonify(request.base_url.replace(self.submit_suffix, self.pickup_suffix)) def pickup(self): result = self.compute_next() @@ -64,7 +67,7 @@ def valid(result): def set_up_processing_server(process_fn): app = Flask(__name__) - stream = RestStreamProcessor(process_fn) + stream = RestStreamProcessor(process_fn, submit_suffix="submit", pickup_suffix="pickup") @app.route("/ready", methods=["GET"]) def ready():