Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

CLI argument for selecting UserClasses to run #2786

Open
wants to merge 4 commits into
base: master
Choose a base branch
from
Open
Show file tree
Hide file tree
Changes from 1 commit
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Next Next commit
add: --run-users cli command
  • Loading branch information
bakhtos committed Jul 2, 2024
commit 96282442f53dad099f02631b3cafcfe4d1117c73
18 changes: 9 additions & 9 deletions locust/argument_parser.py
Original file line number Diff line number Diff line change
Expand Up @@ -201,14 +201,14 @@ def get_empty_argument_parser(add_help=True, default_config_files=DEFAULT_CONFIG
usage=configargparse.SUPPRESS,
description=textwrap.dedent(
"""
Usage: locust [options] [UserClass ...]
Usage: locust [options]
"""
),
epilog="""Examples:

locust -f my_test.py -H https://www.example.com

locust --headless -u 100 -t 20m --processes 4 MyHttpUser AnotherUser
locust --headless -u 100 -t 20m --processes 4 --run-users MyHttpUser AnotherUser

See documentation for more details, including how to set options using a file or environment variables: https://docs.locust.io/en/stable/configuration.html""",
)
Expand Down Expand Up @@ -420,6 +420,12 @@ def setup_parser_arguments(parser):
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",
)
parser.add_argument(
"--run-users",
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think --user-classes is a better name

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This creates an error since here

u 100 # Short form

you have short-form u 100, which gets parsed as --u 100 instead of -u 100 for some reason and so it gives an error because it cannot distinguish between --users and --user-classes if we use such a name

Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yea I figured what might cause an issue. I think --run-users is a bad name though. Can you think of something else?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The best way would have been changing --users to --num_users and then adding --user-classes, but this is not backward-compatible.

nargs="*",
help="List of User classes to be used (available User classes can be listed with --list). LOCUST_USER_CLASSES environment variable can also be used to specify User classes. Default is to use all available User classes",
default=os.environ.get("LOCUST_USER_CLASSES", "").split(),
)

web_ui_group = parser.add_argument_group("Web UI options")
web_ui_group.add_argument(
Expand Down Expand Up @@ -743,13 +749,7 @@ def setup_parser_arguments(parser):
)

user_classes_group = parser.add_argument_group("User classes")
user_classes_group.add_argument(
"user_classes",
nargs="*",
metavar="<UserClass1 UserClass2>",
help="At the end of the command line, you can list User classes to be used (available User classes can be listed with --list). LOCUST_USER_CLASSES environment variable can also be used to specify User classes. Default is to use all available User classes",
default=os.environ.get("LOCUST_USER_CLASSES", "").split(),
)
user_classes_group.add_argument("user_classes", nargs="*")


def get_parser(default_config_files=DEFAULT_CONFIG_FILES) -> LocustArgumentParser:
Expand Down
21 changes: 16 additions & 5 deletions locust/main.py
Original file line number Diff line number Diff line change
Expand Up @@ -13,6 +13,7 @@
import sys
import time
import traceback
import warnings

import gevent

Expand Down Expand Up @@ -310,16 +311,26 @@ def kill_workers(children):
sys.exit(1)

# make sure specified User exists
names = set()
if options.run_users:
if missing := set(options.run_users) - set(user_classes.keys()):
logger.error(f"Unknown User(s): {', '.join(missing)}\n")
sys.exit(1)
else:
names |= set(options.run_users) & set(user_classes.keys())
if options.user_classes:
warnings.warn(
"Specifying users at the end of the command is deprecated. Use --run-users parameter instead",
DeprecationWarning,
)
if missing := set(options.user_classes) - set(user_classes.keys()):
logger.error(f"Unknown User(s): {', '.join(missing)}\n")
sys.exit(1)
else:
names = set(options.user_classes) & set(user_classes.keys())
user_classes = [user_classes[n] for n in names]
else:
# list() call is needed to consume the dict_view object in Python 3
user_classes = list(user_classes.values())
names |= set(options.user_classes) & set(user_classes.keys())
if not (options.run_users or options.user_classes):
names = set(user_classes.keys())
user_classes = [user_classes[n] for n in names]

if os.name != "nt" and not options.master:
try:
Expand Down
3 changes: 2 additions & 1 deletion locust/test/test_main.py
Original file line number Diff line number Diff line change
Expand Up @@ -75,7 +75,7 @@ def test_help_arg(self):
timeout=5,
text=True,
).strip()
self.assertTrue(output.startswith("Usage: locust [options] [UserClass"))
self.assertTrue(output.startswith("Usage: locust [options]"))
self.assertIn("Common options:", output)
self.assertIn("-f <filename>, --locustfile <filename>", output)
self.assertIn("Logging options:", output)
Expand Down Expand Up @@ -1087,6 +1087,7 @@ def t(self):
"5",
"-r",
"10",
"--run-users",
"User2",
"User3",
]
Expand Down