Merge pull request #3026 from ftb-skry/argument-parsing-cleanup

Move some argument parsing/validation from main.py to argument_parser.py and remove deprecated parameter --hatch-rate

locust/argument_parser.py

@@ -6,6 +6,7 @@
 
 import ast
 import atexit
+import json
 import os
 import platform
 import socket
@@ -28,6 +29,7 @@
 import requests
 
 from .util.directory import get_abspaths_in
+from .util.timespan import parse_timespan
 from .util.url import is_url
 
 version = locust.__version__
@@ -369,6 +371,66 @@ def parse_locustfile_option(args=None) -> list[str]:
     return parsed_paths
 
 
+# A hack for setting up an action that raises ArgumentError with configurable error messages.
+# This is meant to be used to immediately block use of deprecated arguments with some helpful messaging.
+
+
+def raise_argument_type_error(err_msg):
+    class ErrorRaisingAction(configargparse.Action):
+        def __call__(self, parser, namespace, values, option_string=None):
+            raise configargparse.ArgumentError(self, err_msg)
+
+    return ErrorRaisingAction
+
+
+# Definitions for some "types" to use with the arguments
+
+
+def timespan(time_str) -> int:
+    try:
+        return parse_timespan(time_str)
+    except ValueError as e:
+        raise configargparse.ArgumentTypeError(str(e))
+
+
+def positive_integer(string) -> int:
+    try:
+        value = int(string)
+    except ValueError:
+        raise configargparse.ArgumentTypeError(f"invalid int value: '{string}'")
+
+    if value < 1:
+        raise configargparse.ArgumentTypeError(
+            f"Invalid --expect-workers argument ({value}), must be a positive number"
+        )
+
+    return value
+
+
+def json_user_config(string):
+    try:
+        if string.endswith(".json"):
+            with open(string) as file:
+                user_config = json.load(file)
+        else:
+            user_config = json.loads(string)
+
+        if not isinstance(user_config, list):
+            user_config = [user_config]
+
+        for config in user_config:
+            if "user_class_name" not in config:
+                raise configargparse.ArgumentTypeError("The user config must specify a user_class_name")
+
+        return user_config
+
+    except json.decoder.JSONDecodeError as e:
+        raise configargparse.ArgumentTypeError(f"The --config-users argument must be a valid JSON string or file: {e}")
+
+    except FileNotFoundError as e:
+        raise configargparse.ArgumentTypeError(str(e))
+
+
 def setup_parser_arguments(parser):
     """
     Setup command-line options
@@ -401,20 +463,13 @@ def setup_parser_arguments(parser):
         help="Rate to spawn users at (users per second). Primarily used together with --headless or --autostart",
         env_var="LOCUST_SPAWN_RATE",
     )
-    parser.add_argument(
-        "--hatch-rate",
-        env_var="LOCUST_HATCH_RATE",
-        metavar="<float>",
-        type=float,
-        default=0,
-        help=configargparse.SUPPRESS,
-    )
     parser.add_argument(
         "-t",
         "--run-time",
         metavar="<time string>",
         help="Stop after the specified amount of time, e.g. (300s, 20m, 3h, 1h30m, etc.). Only used together with --headless or --autostart. Defaults to run forever.",
         env_var="LOCUST_RUN_TIME",
+        type=timespan,
     )
     parser.add_argument(
         "-l",
@@ -425,7 +480,7 @@ def setup_parser_arguments(parser):
     )
     parser.add_argument(
         "--config-users",
-        type=str,
+        type=json_user_config,
         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",
@@ -489,6 +544,9 @@ def setup_parser_arguments(parser):
         default=None,
         help=configargparse.SUPPRESS,
         env_var="LOCUST_WEB_AUTH",
+        action=raise_argument_type_error(
+            "The --web-auth parameters has been replaced with --web-login. See https://docs.locust.io/en/stable/extending-locust.html#authentication for details"
+        ),
     )
     web_ui_group.add_argument(
         "--web-login",
@@ -528,7 +586,8 @@ def setup_parser_arguments(parser):
     web_ui_group.add_argument(
         "--legacy-ui",
         default=False,
-        action="store_true",
+        action=raise_argument_type_error("--legacy-ui is no longer supported, remove the parameter to continue"),
+        nargs=0,
         help=configargparse.SUPPRESS,
         env_var="LOCUST_LEGACY_UI",
     )
@@ -561,7 +620,7 @@ def setup_parser_arguments(parser):
     )
     master_group.add_argument(
         "--expect-workers",
-        type=int,
+        type=positive_integer,
         metavar="<int>",
         default=1,
         help="Delay starting the test until this number of workers have connected (only used in combination with --headless/--autostart).",
@@ -584,7 +643,8 @@ def setup_parser_arguments(parser):
     )
     master_group.add_argument(
         "--expect-slaves",
-        action="store_true",
+        action=raise_argument_type_error("The --expect-slaves parameter has been renamed --expect-workers"),
+        nargs=0,
         help=configargparse.SUPPRESS,
     )
 
@@ -608,7 +668,8 @@ def setup_parser_arguments(parser):
     )
     worker_group.add_argument(
         "--slave",
-        action="store_true",
+        action=raise_argument_type_error("The --slave parameter has been renamed --worker"),
+        nargs=0,
         help=configargparse.SUPPRESS,
     )
     worker_group.add_argument(
@@ -720,6 +781,8 @@ def setup_parser_arguments(parser):
         help="Choose between DEBUG/INFO/WARNING/ERROR/CRITICAL. Default is INFO.",
         metavar="<level>",
         env_var="LOCUST_LOGLEVEL",
+        choices=["DEBUG", "INFO", "WARNING", "ERROR", "CRITICAL"],
+        type=str.upper,
     )
     log_group.add_argument(
         "--logfile",
@@ -764,6 +827,7 @@ def setup_parser_arguments(parser):
         default="0",
         help="Number of seconds to wait for a simulated user to complete any executing task before exiting. Default is to terminate immediately. When running distributed, this only needs to be specified on the master.",
         env_var="LOCUST_STOP_TIMEOUT",
+        type=timespan,
     )
     other_group.add_argument(
         "--equal-weights",

locust/main.py

@@ -7,6 +7,7 @@
 import gc
 import importlib.metadata
 import inspect
+import itertools
 import json
 import logging
 import os
@@ -35,7 +36,6 @@
 )
 from .user.inspectuser import print_task_ratio, print_task_ratio_json
 from .util.load_locustfile import load_locustfile
-from .util.timespan import parse_timespan
 
 # import external plugins if  installed to allow for registering custom arguments etc
 try:
@@ -172,35 +172,13 @@ def is_valid_percentile(parameter):
     if options.headful:
         options.headless = False
 
-    if options.slave or options.expect_slaves:
-        sys.stderr.write("The --slave/--expect-slaves parameters have been renamed --worker/--expect-workers\n")
-        sys.exit(1)
-
-    if options.web_auth:
-        sys.stderr.write(
-            "The --web-auth parameters has been replaced with --web-login. See https://docs.locust.io/en/stable/extending-locust.html#authentication for details\n"
-        )
-        sys.exit(1)
-
     if options.autoquit != -1 and not options.autostart:
         sys.stderr.write("--autoquit is only meaningful in combination with --autostart\n")
         sys.exit(1)
 
-    if options.hatch_rate:
-        sys.stderr.write("--hatch-rate parameter has been renamed --spawn-rate\n")
-        sys.exit(1)
-
-    if options.legacy_ui:
-        sys.stderr.write("--legacy-ui is no longer supported, remove the parameter to continue\n")
-        sys.exit(1)
-
     # setup logging
     if not options.skip_log_setup:
-        if options.loglevel.upper() in ["DEBUG", "INFO", "WARNING", "ERROR", "CRITICAL"]:
-            setup_logging(options.loglevel, options.logfile)
-        else:
-            sys.stderr.write("Invalid --loglevel. Valid values are: DEBUG/INFO/WARNING/ERROR/CRITICAL\n")
-            sys.exit(1)
+        setup_logging(options.loglevel, options.logfile)
 
     children = []
     logger = logging.getLogger(__name__)
@@ -312,13 +290,6 @@ def kill_workers(children):
     if sys.version_info < (3, 10):
         logger.warning("Python 3.9 support is deprecated and will be removed soon")
 
-    if options.stop_timeout:
-        try:
-            options.stop_timeout = parse_timespan(options.stop_timeout)
-        except ValueError:
-            logger.error("Valid --stop-timeout formats are: 20, 20s, 3m, 2h, 1h20m, 3h30m10s, etc.")
-            sys.exit(1)
-
     if options.list_commands:
         print("Available Users:")
         for name in user_classes:
@@ -382,39 +353,17 @@ def kill_workers(children):
     )
 
     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 json.decoder.JSONDecodeError as e:
-                logger.error(f"The --config-users argument must be in valid JSON string or file: {e}")
-                sys.exit(-1)
-            except KeyError as e:
-                logger.error(
-                    f"Error applying user config, probably you tried to specify config for a User not present in your locustfile: {e}"
-                )
-                sys.exit(-1)
-            except Exception as e:
-                logger.exception(e)
-                sys.exit(-1)
+        try:
+            for user_config in itertools.chain(*options.config_users):
+                environment.update_user_class(user_config)
+        except KeyError as e:
+            logger.error(
+                f"Error applying user config, probably you tried to specify config for a User not present in your locustfile: {e}"
+            )
+            sys.exit(-1)
+        except Exception as e:
+            logger.exception(e)
+            sys.exit(-1)
 
     if (
         shape_class
@@ -469,13 +418,6 @@ def ensure_user_class_name(config):
     if options.run_time:
         if options.worker:
             logger.info("--run-time specified for a worker node will be ignored.")
-        try:
-            options.run_time = parse_timespan(options.run_time)
-        except ValueError:
-            logger.error(
-                f"Invalid --run-time argument ({options.run_time}), accepted formats are for example 120, 120s, 2m, 3h, 3h30m10s."
-            )
-            sys.exit(1)
 
     if options.csv_prefix:
         base_csv_file = os.path.basename(options.csv_prefix)

locust/test/test_main.py

@@ -411,8 +411,11 @@ def test_invalid_stop_timeout_string(self):
                 text=True,
             )
             stdout, stderr = proc.communicate()
-            self.assertIn("ERROR/locust.main: Valid --stop-timeout formats are", stderr)
-            self.assertEqual(1, proc.returncode)
+            self.assertIn(
+                "locust: error: argument -s/--stop-timeout: Invalid time span format. Valid formats: 20, 20s, 3m, 2h, 1h20m, 3h30m10s, etc.",
+                stderr,
+            )
+            self.assertEqual(2, proc.returncode)
 
     @unittest.skipIf(os.name == "nt", reason="Signal handling on windows is hard")
     def test_headless_spawn_options_wo_run_time(self):
@@ -1053,7 +1056,7 @@ def t(self):
                     mocked.file_path,
                     "--headless",
                     "--run-time",
-                    "0.5",
+                    "1",
                     "-u",
                     "3",
                 ],

locust/test/test_parser.py

@@ -117,9 +117,9 @@ def test_parse_options(self):
         self.assertEqual("locustfile.py", options.locustfile)
         self.assertEqual(100, options.num_users)
         self.assertEqual(10, options.spawn_rate)
-        self.assertEqual("5m", options.run_time)
+        self.assertEqual(300, options.run_time)
         self.assertTrue(options.reset_stats)
-        self.assertEqual("5", options.stop_timeout)
+        self.assertEqual(5, options.stop_timeout)
         self.assertEqual(["MyUserClass"], options.user_classes)
         # check default arg
         self.assertEqual(8089, options.web_port)
@@ -137,9 +137,9 @@ def test_parse_options_from_env(self):
         self.assertEqual("locustfile.py", options.locustfile)
         self.assertEqual(100, options.num_users)
         self.assertEqual(10, options.spawn_rate)
-        self.assertEqual("5m", options.run_time)
+        self.assertEqual(300, options.run_time)
         self.assertTrue(options.reset_stats)
-        self.assertEqual("5", options.stop_timeout)
+        self.assertEqual(5, options.stop_timeout)
         self.assertEqual(["MyUserClass"], options.user_classes)
         # check default arg
         self.assertEqual(8089, options.web_port)