Skip to content

Commit

Permalink
Add --config-users parser argument
Browse files Browse the repository at this point in the history
  • Loading branch information
andrewbaldwin44 committed Jan 19, 2024
1 parent 190169b commit bfef114
Show file tree
Hide file tree
Showing 2 changed files with 35 additions and 0 deletions.
7 changes: 7 additions & 0 deletions locust/argument_parser.py
Original file line number Diff line number Diff line change
Expand Up @@ -306,6 +306,13 @@ def setup_parser_arguments(parser):
dest="list_commands",
help="Show list of possible User classes and exit",
)
parser.add_argument(
"--config-users",
type=str,
nargs="*",
help="User configuration as a JSON string or file. A list of arguments or an Array of JSON configuration may be provided",
env_var="LOCUST_CONFIG_USERS",
)

web_ui_group = parser.add_argument_group("Web UI options")
web_ui_group.add_argument(
Expand Down
28 changes: 28 additions & 0 deletions locust/main.py
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,7 @@
import errno
import logging
import os
import json
import signal
import sys
import time
Expand Down Expand Up @@ -358,6 +359,33 @@ def kill_workers(children):
available_user_tasks=available_user_tasks,
)

if options.config_users:
for json_user_config in options.config_users:
try:
if json_user_config.endswith(".json"):
with open(json_user_config) as file:
user_config = json.load(file)
else:
user_config = json.loads(json_user_config)

def ensure_user_class_name(config):
if "user_class_name" not in config:
logger.error("The user config must specify a user_class_name")
sys.exit(-1)

if isinstance(user_config, list):
for config in user_config:
ensure_user_class_name(config)

environment.update_user_class(config)
else:
ensure_user_class_name(user_config)

environment.update_user_class(user_config)
except Exception as e:
logger.error(f"The --config-users arugment must be in valid JSON string or file: {e}")
sys.exit(-1)

if (
shape_class
and not shape_class.use_common_options
Expand Down

0 comments on commit bfef114

Please sign in to comment.