diff --git a/locust/argument_parser.py b/locust/argument_parser.py index 827a2f880e..9334b8e8be 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) -> 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="", - type=float, - default=0, - help=configargparse.SUPPRESS, - ) parser.add_argument( "-t", "--run-time", metavar="