Skip to content

Commit

Permalink
Merge pull request #2966 from andrewbaldwin44/feature/base-path
Browse files Browse the repository at this point in the history
Locust Configurable Web Base Path
  • Loading branch information
cyberw authored Nov 5, 2024
2 parents 4996c9e + 409b59f commit caf1baf
Show file tree
Hide file tree
Showing 9 changed files with 65 additions and 44 deletions.
13 changes: 6 additions & 7 deletions examples/web_ui_auth/basic.py
Original file line number Diff line number Diff line change
Expand Up @@ -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",
Expand All @@ -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():
Expand All @@ -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)
13 changes: 6 additions & 7 deletions examples/web_ui_auth/custom_form.py
Original file line number Diff line number Diff line change
Expand Up @@ -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")
Expand All @@ -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",
},
}
Expand All @@ -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:
Expand All @@ -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)
8 changes: 8 additions & 0 deletions locust/argument_parser.py
Original file line number Diff line number Diff line change
Expand Up @@ -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.",
Expand Down
2 changes: 2 additions & 0 deletions locust/env.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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

Expand Down
16 changes: 12 additions & 4 deletions locust/main.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down Expand Up @@ -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,
Expand Down
5 changes: 3 additions & 2 deletions locust/runners.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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()
Expand Down
1 change: 0 additions & 1 deletion locust/test/test_runners.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
6 changes: 3 additions & 3 deletions locust/test/test_web.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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

Expand Down
Loading

0 comments on commit caf1baf

Please sign in to comment.