diff --git a/docs/custom-load-shape.rst b/docs/custom-load-shape.rst index d767b51965..2ddea794d6 100644 --- a/docs/custom-load-shape.rst +++ b/docs/custom-load-shape.rst @@ -78,6 +78,7 @@ Adding the element ``user_classes`` to the return value gives you more detailed {"duration": 30, "users": 50, "spawn_rate": 10, "user_classes": [UserA, UserB]}, {"duration": 60, "users": 100, "spawn_rate": 10, "user_classes": [UserB]}, {"duration": 120, "users": 100, "spawn_rate": 10, "user_classes": [UserA,UserB]}, + ] def tick(self): run_time = self.get_run_time() @@ -93,3 +94,30 @@ Adding the element ``user_classes`` to the return value gives you more detailed return None This shape would create create in the first 10 seconds 10 User of ``UserA``. In the next twenty seconds 40 of type ``UserA / UserB`` and this continues until the stages end. + + +.. _use-common-options: + +Reusing command line parameters in custom shapes +------------------------------------------------ + +By default, using a custom shape will disable default run paramaters (in both the CLI and the Web UI): +- `--run-time` (providing this one with a custom shape will make locust to bail out) +- `--spawn-rate` +- `--users` + + +If you need one or all of those parameters, you can force locust to accept them by setting the `use_common_options` attribute to `True`: + + +.. code-block:: python + + class MyCustomShape(LoadTestShape): + + use_common_options = True + + def tick(self): + expected_run_time = self.runner.environment.parsed_options.run_time + # Do something with this expected run time + ... + return None diff --git a/locust/main.py b/locust/main.py index 4591c4e2e7..1c27c093f8 100644 --- a/locust/main.py +++ b/locust/main.py @@ -38,6 +38,15 @@ version = locust.__version__ +# Options to ignore when using a custom shape class without `use_common_options=True` +# See: https://docs.locust.io/en/stable/custom-load-shape.html#use-common-options +COMMON_OPTIONS = { + "num_users": "users", + "spawn_rate": "spawn-rate", + "run_time": "run-time", +} + + def create_environment( user_classes, options, @@ -214,10 +223,17 @@ def is_valid_percentile(parameter): available_shape_classes=available_shape_classes, ) - if shape_class and (options.num_users or options.spawn_rate): + if ( + shape_class + and not shape_class.use_common_options + and any(getattr(options, opt, None) for opt in COMMON_OPTIONS) + ): logger.warning( - "The specified locustfile contains a shape class but a conflicting argument was specified: users or spawn-rate. Ignoring arguments" + "--run-time, --users or --spawn-rate have no impact on LoadShapes unless the shape class explicitly reads them. " + "See: docs.locust.io/en/stable/custom-load-shape.html#use-common-options" ) + ignored = [f"--{arg}" for opt, arg in COMMON_OPTIONS.items() if getattr(options, opt, None)] + logger.warning(f"The following option(s) will be ignored: {', '.join(ignored)}") if options.show_task_ratio: print("\n Task ratio per User class") @@ -381,9 +397,6 @@ def start_automatic_run(): # start the test if environment.shape_class: - if options.run_time: - sys.stderr.write("It makes no sense to combine --run-time and LoadShapes. Bailing out.\n") - sys.exit(1) try: environment.runner.start_shape() environment.runner.shape_greenlet.join() diff --git a/locust/runners.py b/locust/runners.py index 78b856fcc2..f6250303de 100644 --- a/locust/runners.py +++ b/locust/runners.py @@ -329,7 +329,7 @@ def start_shape(self) -> None: logger.info("There is an ongoing shape test running. Editing is disabled") return - logger.info("Shape test starting. User count and spawn rate are ignored for this type of load test") + logger.info("Shape test starting.") self.update_state(STATE_INIT) self.shape_greenlet = self.greenlet.spawn(self.shape_worker) self.shape_greenlet.link_exception(greenlet_exception_handler) diff --git a/locust/shape.py b/locust/shape.py index afdeccd8fd..77417d13a9 100644 --- a/locust/shape.py +++ b/locust/shape.py @@ -28,6 +28,8 @@ class LoadTestShape(metaclass=LoadTestShapeMeta): abstract: ClassVar[bool] = True + use_common_options: ClassVar[bool] = False + def __init__(self): self.start_time = time.perf_counter() diff --git a/locust/templates/index.html b/locust/templates/index.html index 3e66fb3d5b..2abb59897d 100644 --- a/locust/templates/index.html +++ b/locust/templates/index.html @@ -65,7 +65,7 @@