diff --git a/locust/argument_parser.py b/locust/argument_parser.py index c4ee468fea..96510c768b 100644 --- a/locust/argument_parser.py +++ b/locust/argument_parser.py @@ -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 --user-classes 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""", ) @@ -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( + "--user-classes", + 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( @@ -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="", - 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("old_user_classes", nargs="*") def get_parser(default_config_files=DEFAULT_CONFIG_FILES) -> LocustArgumentParser: diff --git a/locust/main.py b/locust/main.py index 317db376c6..f01f3f82f4 100644 --- a/locust/main.py +++ b/locust/main.py @@ -310,16 +310,27 @@ def kill_workers(children): sys.exit(1) # make sure specified User exists + names = set() if options.user_classes: 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) + if options.old_user_classes: + # TODO deprecate in future release + # warnings.warn( + # "Specifying users at the end of the command is deprecated. Use --run-users parameter instead", + # DeprecationWarning, + # ) + if missing := set(options.old_user_classes) - set(user_classes.keys()): + logger.error(f"Unknown User(s): {', '.join(missing)}\n") + sys.exit(1) + else: + names |= set(options.old_user_classes) + if not (options.old_user_classes 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: diff --git a/locust/test/test_load_locustfile.py b/locust/test/test_load_locustfile.py index d3058deeb4..cd5f4cef8b 100644 --- a/locust/test/test_load_locustfile.py +++ b/locust/test/test_load_locustfile.py @@ -166,9 +166,8 @@ def test_specify_config_file(self): textwrap.dedent( """ host = localhost # With "=" - u 100 # Short form - spawn-rate 5 # long form - # boolean + users 100 + spawn-rate 5 headless # (for some reason an inline comment makes boolean values fail in configargparse nowadays) """ diff --git a/locust/test/test_main.py b/locust/test/test_main.py index f50625257a..5d6cfdcb95 100644 --- a/locust/test/test_main.py +++ b/locust/test/test_main.py @@ -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 , --locustfile ", output) self.assertIn("Logging options:", output) @@ -1087,6 +1087,7 @@ def t(self): "5", "-r", "10", + "--user-classes", "User2", "User3", ] @@ -2202,6 +2203,7 @@ def test_processes_error_doesnt_blow_up_completely(self): "4", "-L", "DEBUG", + "--user-classes", "UserThatDoesntExist", ], stdout=PIPE, diff --git a/locust/test/test_parser.py b/locust/test/test_parser.py index a76cc9d7fb..9bc4a6da5c 100644 --- a/locust/test/test_parser.py +++ b/locust/test/test_parser.py @@ -111,6 +111,7 @@ def test_parse_options(self): "--reset-stats", "--stop-timeout", "5", + "--user-classes", "MyUserClass", ] ) @@ -167,6 +168,7 @@ def test_parse_locustfile(self): "--reset-stats", "--stop-timeout", "5", + "--user-classes", "MyUserClass", ] )