Merge pull request locustio#2966 from andrewbaldwin44/feature/base-path

Locust Configurable Web Base Path

examples/web_ui_auth/basic.py

@@ -29,22 +29,21 @@ def get_id(self):
         return self.username
 
 
-auth_blueprint = Blueprint("auth", "web_ui_auth")
-
-
 def load_user(username):
     return AuthUser(username)
 
 
 @events.init.add_listener
 def locust_init(environment, **_kwargs):
     if environment.web_ui:
+        auth_blueprint = Blueprint("auth", "web_ui_auth", url_prefix=environment.parsed_options.web_base_path)
+
         environment.web_ui.login_manager.user_loader(load_user)
 
         environment.web_ui.app.config["SECRET_KEY"] = os.getenv("FLASK_SECRET_KEY")
 
         environment.web_ui.auth_args = {
-            "username_password_callback": "/login_submit",
+            "username_password_callback": f"{environment.parsed_options.web_base_path}/login_submit",
             "auth_providers": [
                 {
                     "label": "Github",
@@ -61,7 +60,7 @@ def google_login():
             session["username"] = username
             login_user(AuthUser("username"))
 
-            return redirect(url_for("index"))
+            return redirect(url_for("locust.index"))
 
         @auth_blueprint.route("/login_submit", methods=["POST"])
         def login_submit():
@@ -72,10 +71,10 @@ def login_submit():
             if password:
                 login_user(AuthUser(username))
 
-                return redirect(url_for("index"))
+                return redirect(url_for("locust.index"))
 
             session["auth_error"] = "Invalid username or password"
 
-            return redirect(url_for("login"))
+            return redirect(url_for("locust.login"))
 
         environment.web_ui.app.register_blueprint(auth_blueprint)

examples/web_ui_auth/custom_form.py

@@ -34,16 +34,15 @@ def get_id(self):
         return self.username
 
 
-auth_blueprint = Blueprint("auth", "web_ui_auth")
-
-
 def load_user(user_id):
     return AuthUser(user_id)
 
 
 @events.init.add_listener
 def locust_init(environment, **_kwargs):
     if environment.web_ui:
+        auth_blueprint = Blueprint("auth", "web_ui_auth", url_prefix=environment.parsed_options.web_base_path)
+
         environment.web_ui.login_manager.user_loader(load_user)
 
         environment.web_ui.app.config["SECRET_KEY"] = os.getenv("FLASK_SECRET_KEY")
@@ -70,7 +69,7 @@ def locust_init(environment, **_kwargs):
                         "is_secret": True,
                     },
                 ],
-                "callback_url": "/login_submit",
+                "callback_url": f"{environment.parsed_options.web_base_path}/login_submit",
                 "submit_button_text": "Submit",
             },
         }
@@ -86,7 +85,7 @@ def login_submit():
             if password != confirm_password:
                 session["auth_error"] = "Passwords do not match!"
 
-                return redirect(url_for("login"))
+                return redirect(url_for("locust.login"))
 
             # Implement real password verification here
             if password:
@@ -98,10 +97,10 @@ def login_submit():
 
                 login_user(AuthUser(username))
 
-                return redirect(url_for("index"))
+                return redirect(url_for("locust.index"))
 
             session["auth_error"] = "Invalid username or password"
 
-            return redirect(url_for("login"))
+            return redirect(url_for("locust.login"))
 
         environment.web_ui.app.register_blueprint(auth_blueprint)

locust/argument_parser.py

@@ -612,6 +612,14 @@ def setup_parser_arguments(parser):
         env_var="LOCUST_MASTER_NODE_PORT",
     )
 
+    web_ui_group.add_argument(
+        "--web-base-path",
+        type=str,
+        default="",
+        help="Base path for the web interface (e.g., '/locust'). Default is empty (root path).",
+        env_var="LOCUST_web_base_path",
+    )
+
     tag_group = parser.add_argument_group(
         "Tag options",
         "Locust tasks can be tagged using the @tag decorator. These options let specify which tasks to include or exclude during a test.",

locust/env.py

@@ -165,6 +165,7 @@ def create_web_ui(
         self,
         host="",
         port=8089,
+        web_base_path: str | None = None,
         web_login: bool = False,
         tls_cert: str | None = None,
         tls_key: str | None = None,
@@ -199,6 +200,7 @@ def create_web_ui(
             delayed_start=delayed_start,
             userclass_picker_is_active=userclass_picker_is_active,
             build_path=build_path,
+            web_base_path=web_base_path,
         )
         return self.web_ui
 

locust/main.py

@@ -454,7 +454,9 @@ def ensure_user_class_name(config):
     elif options.worker:
         try:
             runner = environment.create_worker_runner(options.master_host, options.master_port)
-            logger.debug("Connected to locust master: %s:%s", options.master_host, options.master_port)
+            logger.debug(
+                "Connected to locust master: %s:%s%s", options.master_host, options.master_port, options.web_base_path
+            )
         except OSError as e:
             logger.error("Failed to connect to the Locust master: %s", e)
             sys.exit(-1)
@@ -490,26 +492,32 @@ def ensure_user_class_name(config):
     if not options.headless and not options.worker:
         protocol = "https" if options.tls_cert and options.tls_key else "http"
 
+        if options.web_base_path and options.web_base_path[0] != "/":
+            logger.error(
+                f"Invalid format for --web-base-path argument ({options.web_base_path}): the url path must start with a slash."
+            )
+            sys.exit(1)
         if options.web_host == "*":
             # special check for "*" so that we're consistent with --master-bind-host
             web_host = ""
         else:
             web_host = options.web_host
         if web_host:
-            logger.info(f"Starting web interface at {protocol}://{web_host}:{options.web_port}")
+            logger.info(f"Starting web interface at {protocol}://{web_host}:{options.web_port}{options.web_base_path}")
         if options.web_host_display_name:
             logger.info(f"Starting web interface at {options.web_host_display_name}")
         else:
             if os.name == "nt":
                 logger.info(
-                    f"Starting web interface at {protocol}://localhost:{options.web_port} (accepting connections from all network interfaces)"
+                    f"Starting web interface at {protocol}://localhost:{options.web_port}{options.web_base_path} (accepting connections from all network interfaces)"
                 )
             else:
-                logger.info(f"Starting web interface at {protocol}://0.0.0.0:{options.web_port}")
+                logger.info(f"Starting web interface at {protocol}://0.0.0.0:{options.web_port}{options.web_base_path}")
 
         web_ui = environment.create_web_ui(
             host=web_host,
             port=options.web_port,
+            web_base_path=options.web_base_path,
             web_login=options.web_login,
             tls_cert=options.tls_cert,
             tls_key=options.tls_key,

locust/runners.py

@@ -1216,6 +1216,7 @@ def __init__(self, environment: Environment, master_host: str, master_port: int)
         self.client_id = socket.gethostname() + "_" + uuid4().hex
         self.master_host = master_host
         self.master_port = master_port
+        self.web_base_path = environment.parsed_options.web_base_path if environment.parsed_options else ""
         self.logs: list[str] = []
         self.worker_cpu_warning_emitted = False
         self._users_dispatcher: UsersDispatcher | None = None
@@ -1475,11 +1476,11 @@ def connect_to_master(self):
         if not success:
             if self.retry < 3:
                 logger.debug(
-                    f"Failed to connect to master {self.master_host}:{self.master_port}, retry {self.retry}/{CONNECT_RETRY_COUNT}."
+                    f"Failed to connect to master {self.master_host}:{self.master_port}{self.web_base_path}, retry {self.retry}/{CONNECT_RETRY_COUNT}."
                 )
             else:
                 logger.warning(
-                    f"Failed to connect to master {self.master_host}:{self.master_port}, retry {self.retry}/{CONNECT_RETRY_COUNT}."
+                    f"Failed to connect to master {self.master_host}:{self.master_port}{self.web_base_path}, retry {self.retry}/{CONNECT_RETRY_COUNT}."
                 )
             if self.retry > CONNECT_RETRY_COUNT:
                 raise ConnectionError()

locust/test/test_runners.py

@@ -73,7 +73,6 @@ def send(self, message):
             self.outbox.append(message)
 
         def send_to_client(self, message):
-            print(message)
             self.outbox.append(message)
 
         @classmethod

locust/test/test_web.py

@@ -55,7 +55,7 @@ def setUp(self):
         self.stats = self.environment.stats
 
         self.web_ui = self.environment.create_web_ui("127.0.0.1", 0)
-        self.web_ui.app.view_functions["request_stats"].clear_cache()
+        self.web_ui.app.view_functions["locust.request_stats"].clear_cache()
         gevent.sleep(0.01)
         self.web_port = self.web_ui.server.server_port
 
@@ -148,7 +148,7 @@ def test_stats_cache(self):
         data = json.loads(requests.get("http://127.0.0.1:%i/stats/requests" % self.web_port).text)
         self.assertEqual(2, len(data["stats"]))  # old value should be cached now
 
-        self.web_ui.app.view_functions["request_stats"].clear_cache()
+        self.web_ui.app.view_functions["locust.request_stats"].clear_cache()
 
         data = json.loads(requests.get("http://127.0.0.1:%i/stats/requests" % self.web_port).text)
         self.assertEqual(3, len(data["stats"]))  # this should no longer be cached
@@ -1208,7 +1208,7 @@ def setUp(self):
             self.environment, stats.PERCENTILES_TO_REPORT, self.STATS_BASE_NAME, full_history=True
         )
         self.web_ui = self.environment.create_web_ui("127.0.0.1", 0, stats_csv_writer=self.stats_csv_writer)
-        self.web_ui.app.view_functions["request_stats"].clear_cache()
+        self.web_ui.app.view_functions["locust.request_stats"].clear_cache()
         gevent.sleep(0.01)
         self.web_port = self.web_ui.server.server_port