Skip to content

Commit

Permalink
Update to use Blueprint
Browse files Browse the repository at this point in the history
  • Loading branch information
@andrewbaldwin44
andrewbaldwin44 committed Nov 4, 2024
1 parent b8fe5fe commit 87d578f
Showing 1 changed file with 21 additions and 19 deletions.
40 changes: 21 additions & 19 deletions locust/web.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,7 @@

import gevent
from flask import (
Blueprint,
Flask,
Response,
jsonify,
Expand Down Expand Up @@ -120,7 +121,7 @@ def __init__(
environment: Environment,
host: str,
port: int,
base_path: str,
base_path: str | None = None,
web_login: bool = False,
tls_cert: str | None = None,
tls_key: str | None = None,
Expand All @@ -147,7 +148,6 @@ def __init__(
self.environment = environment
self.host = host
self.port = port
self.base_path = base_path
self.tls_cert = tls_cert
self.tls_key = tls_key
self.userclass_picker_is_active = userclass_picker_is_active
Expand All @@ -163,6 +163,8 @@ def __init__(
self.auth_args = {}
self.app.template_folder = build_path or DEFAULT_BUILD_PATH
self.app.static_url_path = "/assets/"

app_blueprint = Blueprint("locust", __name__, url_prefix=base_path)
# ensures static js files work on Windows
mimetypes.add_type("application/javascript", ".js")
if self.web_login:
Expand All @@ -175,7 +177,7 @@ def __init__(
if not delayed_start:
self.start()

@app.errorhandler(Exception)
@app_blueprint.errorhandler(Exception)
def handle_exception(error):
error_message = str(error)
error_code = getattr(error, "code", 500)
Expand All @@ -185,7 +187,7 @@ def handle_exception(error):
)
return make_response(error_message, error_code)

@app.route(f"{self.base_path}/assets/<path:path>")
@app_blueprint.route("/assets/<path:path>")
def send_assets(path):
directory = (
os.path.join(self.app.template_folder, "assets")
Expand All @@ -195,7 +197,7 @@ def send_assets(path):

return send_from_directory(directory, path)

@app.route(f"{self.base_path}/")
@app_blueprint.route("/")
@self.auth_required_if_enabled
def index() -> str | Response:
if not environment.runner:
Expand All @@ -204,7 +206,7 @@ def index() -> str | Response:

return render_template("index.html", template_args=self.template_args)

@app.route(f"{self.base_path}/swarm", methods=["POST"])
@app_blueprint.route("/swarm", methods=["POST"])
@self.auth_required_if_enabled
def swarm() -> Response:
assert request.method == "POST"
Expand Down Expand Up @@ -318,7 +320,7 @@ def swarm() -> Response:
else:
return jsonify({"success": False, "message": "No runner", "host": environment.host})

@app.route(f"{self.base_path}/stop")
@app_blueprint.route("/stop")
@self.auth_required_if_enabled
def stop() -> Response:
if self._swarm_greenlet is not None:
Expand All @@ -328,7 +330,7 @@ def stop() -> Response:
environment.runner.stop()
return jsonify({"success": True, "message": "Test stopped"})

@app.route(f"{self.base_path}/stats/reset")
@app_blueprint.route("/stats/reset")
@self.auth_required_if_enabled
def reset_stats() -> str:
environment.events.reset_stats.fire()
Expand All @@ -337,7 +339,7 @@ def reset_stats() -> str:
environment.runner.exceptions = {}
return "ok"

@app.route(f"{self.base_path}/stats/report")
@app_blueprint.route("/stats/report")
@self.auth_required_if_enabled
def stats_report() -> Response:
theme = request.args.get("theme", "")
Expand Down Expand Up @@ -375,15 +377,15 @@ def _download_csv_response(csv_data: str, filename_prefix: str) -> Response:
)
return response

@app.route(f"{self.base_path}/stats/requests/csv")
@app_blueprint.route("/stats/requests/csv")
@self.auth_required_if_enabled
def request_stats_csv() -> Response:
data = StringIO()
writer = csv.writer(data)
self.stats_csv_writer.requests_csv(writer)
return _download_csv_response(data.getvalue(), "requests")

@app.route(f"{self.base_path}/stats/requests_full_history/csv")
@app_blueprint.route("/stats/requests_full_history/csv")
@self.auth_required_if_enabled
def request_stats_full_history_csv() -> Response:
options = self.environment.parsed_options
Expand All @@ -401,15 +403,15 @@ def request_stats_full_history_csv() -> Response:

return make_response("Error: Server was not started with option to generate full history.", 404)

@app.route(f"{self.base_path}/stats/failures/csv")
@app_blueprint.route("/stats/failures/csv")
@self.auth_required_if_enabled
def failures_stats_csv() -> Response:
data = StringIO()
writer = csv.writer(data)
self.stats_csv_writer.failures_csv(writer)
return _download_csv_response(data.getvalue(), "failures")

@app.route(f"{self.base_path}/stats/requests")
@app_blueprint.route("/stats/requests")
@self.auth_required_if_enabled
@memoize(timeout=DEFAULT_CACHE_TIME, dynamic_timeout=True)
def request_stats() -> Response:
Expand Down Expand Up @@ -479,7 +481,7 @@ def request_stats() -> Response:

return jsonify(report)

@app.route(f"{self.base_path}/exceptions")
@app_blueprint.route("/exceptions")
@self.auth_required_if_enabled
def exceptions() -> Response:
return jsonify(
Expand All @@ -496,15 +498,15 @@ def exceptions() -> Response:
}
)

@app.route(f"{self.base_path}/exceptions/csv")
@app_blueprint.route("/exceptions/csv")
@self.auth_required_if_enabled
def exceptions_csv() -> Response:
data = StringIO()
writer = csv.writer(data)
self.stats_csv_writer.exceptions_csv(writer)
return _download_csv_response(data.getvalue(), "exceptions")

@app.route(f"{self.base_path}/tasks")
@app_blueprint.route("/tasks")
@self.auth_required_if_enabled
def tasks() -> dict[str, dict[str, dict[str, float]]]:
runner = self.environment.runner
Expand All @@ -524,12 +526,12 @@ def tasks() -> dict[str, dict[str, dict[str, float]]]:
}
return task_data

@app.route(f"{self.base_path}/logs")
@app_blueprint.route("/logs")
@self.auth_required_if_enabled
def logs():
return jsonify({"master": get_logs(), "workers": self.environment.worker_logs})

@app.route(f"{self.base_path}/login")
@app_blueprint.route("/login")
def login():
if not self.web_login:
return redirect(url_for("index"))
Expand All @@ -542,7 +544,7 @@ def login():
auth_args=self.auth_args,
)

@app.route(f"{self.base_path}/user", methods=["POST"])
@app_blueprint.route("/user", methods=["POST"])
def update_user():
assert request.method == "POST"

Expand Down

0 comments on commit 87d578f

Please sign in to comment.