diff --git a/locust/argument_parser.py b/locust/argument_parser.py index 827a2f880e..13c1a47613 100644 --- a/locust/argument_parser.py +++ b/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): + try: + return parse_timespan(time_str) + except ValueError as e: + raise configargparse.ArgumentTypeError(str(e)) + + +def positive_integer(string): + 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 @@ -408,6 +470,7 @@ def setup_parser_arguments(parser): type=float, default=0, help=configargparse.SUPPRESS, + action=raise_argument_type_error("--hatch-rate parameter has been renamed --spawn-rate"), ) parser.add_argument( "-t", @@ -415,6 +478,7 @@ def setup_parser_arguments(parser): metavar="