From 196e86bab861ffe8b3248ebb51263a5581d3cd39 Mon Sep 17 00:00:00 2001 From: "Axel H." Date: Tue, 5 Sep 2023 22:09:56 +0200 Subject: [PATCH] refactor(shape): allows custom shapes to reuse run-time, spawn-rate and users parameters --- docs/custom-load-shape.rst | 28 +++++++++++++- locust/main.py | 2 +- locust/runners.py | 4 +- locust/shape.py | 4 +- locust/test/test_main.py | 77 ++++++++++++++++++++++++++++++++++++++ locust/test/test_web.py | 48 ++++++++++++++++++++++++ locust/web.py | 5 ++- 7 files changed, 163 insertions(+), 5 deletions(-) diff --git a/docs/custom-load-shape.rst b/docs/custom-load-shape.rst index 2bb013c0a2..68712e8509 100644 --- a/docs/custom-load-shape.rst +++ b/docs/custom-load-shape.rst @@ -66,6 +66,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() @@ -80,4 +81,29 @@ 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. \ No newline at end of file +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. + + +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 `reuse_parameters` attribute to `True`: + + +.. code-block:: python + + class MyCustomShape(LoadTestShape): + + reuse_parameters = 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..b4bed3257b 100644 --- a/locust/main.py +++ b/locust/main.py @@ -381,7 +381,7 @@ def start_automatic_run(): # start the test if environment.shape_class: - if options.run_time: + if options.run_time and not environment.shape_class.reuse_parameters: sys.stderr.write("It makes no sense to combine --run-time and LoadShapes. Bailing out.\n") sys.exit(1) try: diff --git a/locust/runners.py b/locust/runners.py index 78b856fcc2..ce25b26026 100644 --- a/locust/runners.py +++ b/locust/runners.py @@ -329,7 +329,9 @@ 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.") + if self.environment.shape_class and not self.environment.shape_class.reuse_parameters: + logger.info("User count and spawn rate are ignored for this type of load test.") 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 5e09036100..88556ab2e0 100644 --- a/locust/shape.py +++ b/locust/shape.py @@ -1,6 +1,6 @@ from __future__ import annotations import time -from typing import Optional, Tuple, List, Type +from typing import ClassVar, Optional, Tuple, List, Type from abc import ABC, abstractmethod from . import User @@ -15,6 +15,8 @@ class LoadTestShape(ABC): runner: Optional[Runner] = None """Reference to the :class:`Runner ` instance""" + reuse_parameters: ClassVar[bool] = False + def __init__(self): self.start_time = time.perf_counter() diff --git a/locust/test/test_main.py b/locust/test/test_main.py index 6048f3d15a..9bd6209eb7 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,76 @@ 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: + proc = subprocess.Popen( + [ + "locust", + "-f", + mocked.file_path, + "--run-time=1s", + "--headless", + "--exit-code-on-error", + "0", + ], + stdout=PIPE, + stderr=PIPE, + text=True, + ) + gevent.sleep(1) + _, stderr = proc.communicate() + self.assertIn("Bailing out", stderr) + self.assertEqual(1, proc.returncode) + + 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"]) + self.assertIn("Shape test starting.", out.stderr) + self.assertIn("User count and spawn rate are ignored for this type of load test", out.stderr) + + def test_shape_class_with_reuse_parameters(self): + content = MOCK_LOCUSTFILE_CONTENT + textwrap.dedent( + """ + from locust import LoadTestShape + + class TestShape(LoadTestShape): + reuse_parameters = 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", + "--headless", + "--exit-code-on-error=0", + ] + ) + self.assertIn("Shape test starting.", out.stderr) + self.assertNotIn("User count and spawn rate are ignored for this type of load test", 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..d44225ff01 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_reuse_parameters_keep_num_users_and_spawn_rate(self): + class TestShape(LoadTestShape): + reuse_parameters = 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..75d30c5db6 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, + "is_shape": ( + self.environment.shape_class + and not (self.userclass_picker_is_active or self.environment.shape_class.reuse_parameters) + ), "stats_history_enabled": options and options.stats_history_enabled, "tasks": dumps({}), "extra_options": extra_options,