Skip to content

Commit

Permalink
Add --user-config parser argument
Browse files Browse the repository at this point in the history
  • Loading branch information
andrewbaldwin44 committed Jan 18, 2024
1 parent 0561adb commit d3712f5
Show file tree
Hide file tree
Showing 2 changed files with 34 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(
"--user-config",
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_USER_CONFIG",
)

web_ui_group = parser.add_argument_group("Web UI options")
web_ui_group.add_argument(
Expand Down
27 changes: 27 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,32 @@ def kill_workers(children):
available_user_tasks=available_user_tasks,
)

if options.user_config:
for json_user_config in options.user_config:
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)

if isinstance(user_config, list):
for config in user_config:
if "user_class_name" not in config:
logger.error("The user config must specify a user_class_name")
sys.exit(-1)

environment.update_user_class(config)
else:
if "user_class_name" not in user_config:
logger.error("The user config must specify a user_class_name")
sys.exit(-1)

environment.update_user_class(user_config)
except Exception as e:
logger.error(f"The --user-config 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 d3712f5

Please sign in to comment.