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 @@

Start new load test

{% for user in available_user_classes %} {% endfor %} -
+
range
{% endif %} - {% if is_shape %} + {% if hide_common_options %}

@@ -144,7 +144,7 @@

Start new load test

Edit running load test

- {% if is_shape %} + {% if hide_common_options %}
@@ -155,7 +155,7 @@

Edit running load test


{% endif %} - {% if is_shape %} + {% if hide_common_options %} {% else %} diff --git a/locust/test/test_main.py b/locust/test/test_main.py index 6048f3d15a..f3f59267d6 100644 --- a/locust/test/test_main.py +++ b/locust/test/test_main.py @@ -1,3 +1,5 @@ +from __future__ import annotations + import json import os import platform @@ -62,6 +64,11 @@ def test_help_arg(self): self.assertIn("Logging options:", output) self.assertIn("--skip-log-setup Disable Locust's logging setup.", output) + def assert_run(self, cmd: list[str], timeout: int = 5) -> subprocess.CompletedProcess[str]: + out = subprocess.run(cmd, capture_output=True, text=True, timeout=timeout) + self.assertEqual(0, out.returncode, f"locust run failed with exit code {out.returncode}:\n{out.stderr}") + return out + class StandaloneIntegrationTests(ProcessIntegrationTest): def test_custom_arguments(self): @@ -1144,6 +1151,86 @@ def tick(self): self.assertIn("Duplicate shape classes: TestShape", stderr) self.assertEqual(1, proc.returncode) + def test_error_when_providing_both_run_time_and_a_shape_class(self): + content = MOCK_LOCUSTFILE_CONTENT + textwrap.dedent( + """ + from locust import LoadTestShape + class TestShape(LoadTestShape): + def tick(self): + return None + """ + ) + with mock_locustfile(content=content) as mocked: + out = self.assert_run( + [ + "locust", + "-f", + mocked.file_path, + "--run-time=1s", + "--headless", + "--exit-code-on-error", + "0", + ] + ) + + self.assertIn("--run-time, --users or --spawn-rate have no impact on LoadShapes", out.stderr) + self.assertIn("The following option(s) will be ignored: --run-time", out.stderr) + + def test_shape_class_log_disabled_parameters(self): + content = MOCK_LOCUSTFILE_CONTENT + textwrap.dedent( + """ + from locust import LoadTestShape + + class TestShape(LoadTestShape): + def tick(self): + return None + """ + ) + with mock_locustfile(content=content) as mocked: + out = self.assert_run( + [ + "locust", + "--headless", + "-f", + mocked.file_path, + "--exit-code-on-error=0", + "--users=1", + "--spawn-rate=1", + ] + ) + self.assertIn("Shape test starting.", out.stderr) + self.assertIn("--run-time, --users or --spawn-rate have no impact on LoadShapes", out.stderr) + self.assertIn("The following option(s) will be ignored: --users, --spawn-rate", out.stderr) + + def test_shape_class_with_use_common_options(self): + content = MOCK_LOCUSTFILE_CONTENT + textwrap.dedent( + """ + from locust import LoadTestShape + + class TestShape(LoadTestShape): + use_common_options = True + + def tick(self): + return None + """ + ) + with mock_locustfile(content=content) as mocked: + out = self.assert_run( + [ + "locust", + "-f", + mocked.file_path, + "--run-time=1s", + "--users=1", + "--spawn-rate=1", + "--headless", + "--exit-code-on-error=0", + ] + ) + self.assertIn("Shape test starting.", out.stderr) + self.assertNotIn("--run-time, --users or --spawn-rate have no impact on LoadShapes", out.stderr) + self.assertNotIn("The following option(s) will be ignored:", out.stderr) + def test_error_when_locustfiles_directory_is_empty(self): with TemporaryDirectory() as temp_dir: proc = subprocess.Popen(["locust", "-f", temp_dir], stdout=PIPE, stderr=PIPE, text=True) diff --git a/locust/test/test_web.py b/locust/test/test_web.py index ff393727f1..dafd9addb2 100644 --- a/locust/test/test_web.py +++ b/locust/test/test_web.py @@ -945,6 +945,54 @@ def test_report_exceptions(self): isinstance(next(iter(self.runner.exceptions.values()))["nodes"], set), "exception object has been mutated" ) + def test_custom_shape_deactivate_num_users_and_spawn_rate(self): + class TestShape(LoadTestShape): + def tick(self): + return None + + self.environment.shape_class = TestShape + + response = requests.get("http://127.0.0.1:%i/" % self.web_port) + self.assertEqual(200, response.status_code) + + # regex to match the intended select tag with id from the custom argument + re_disabled_user_count = re.compile( + r"]*id=\"(new_)?user_count\"[^>]*disabled=\"disabled\"[^>]*>", flags=re.I + ) + self.assertRegex(response.text, re_disabled_user_count) + + re_disabled_spawn_rate = re.compile( + r"]*id=\"(new_)?spawn_rate\"[^>]*disabled=\"disabled\"[^>]*>", flags=re.I + ) + self.assertRegex(response.text, re_disabled_spawn_rate) + + def test_custom_shape_with_use_common_options_keep_num_users_and_spawn_rate(self): + class TestShape(LoadTestShape): + use_common_options = True + + def tick(self): + return None + + self.environment.shape_class = TestShape + + response = requests.get("http://127.0.0.1:%i/" % self.web_port) + self.assertEqual(200, response.status_code) + + # regex to match the intended select tag with id from the custom argument + re_user_count = re.compile(r"]*id=\"(new_)?user_count\"[^>]*>", flags=re.I) + re_disabled_user_count = re.compile( + r"]*id=\"(new_)?user_count\"[^>]*disabled=\"disabled\"[^>]*>", flags=re.I + ) + self.assertRegex(response.text, re_user_count) + self.assertNotRegex(response.text, re_disabled_user_count) + + re_spawn_rate = re.compile(r"]*id=\"(new_)?spawn_rate\"[^>]*>", flags=re.I) + re_disabled_spawn_rate = re.compile( + r"]*id=\"(new_)?spawn_rate\"[^>]*disabled=\"disabled\"[^>]*>", flags=re.I + ) + self.assertRegex(response.text, re_spawn_rate) + self.assertNotRegex(response.text, re_disabled_spawn_rate) + class TestWebUIAuth(LocustTestCase): def setUp(self): diff --git a/locust/web.py b/locust/web.py index 590fbe8b8f..4cecd464c4 100644 --- a/locust/web.py +++ b/locust/web.py @@ -553,7 +553,10 @@ def update_template_args(self): "num_users": options and options.num_users, "spawn_rate": options and options.spawn_rate, "worker_count": worker_count, - "is_shape": self.environment.shape_class and not self.userclass_picker_is_active, + "hide_common_options": ( + self.environment.shape_class + and not (self.userclass_picker_is_active or self.environment.shape_class.use_common_options) + ), "stats_history_enabled": options and options.stats_history_enabled, "tasks": dumps({}), "extra_options": extra_options,