refactor(shape): allows custom shapes to reuse run-time, spawn-rate a…

…nd users parameters

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.
+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

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:

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)

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 <locust.runners.Runner>` instance"""
 
+    reuse_parameters: ClassVar[bool] = False
+
     def __init__(self):
         self.start_time = time.perf_counter()
 

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)

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"<input[^>]*id=\"(new_)?user_count\"[^>]*disabled=\"disabled\"[^>]*>", flags=re.I
+        )
+        self.assertRegex(response.text, re_disabled_user_count)
+
+        re_disabled_spawn_rate = re.compile(
+            r"<input[^>]*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"<input[^>]*id=\"(new_)?user_count\"[^>]*>", flags=re.I)
+        re_disabled_user_count = re.compile(
+            r"<input[^>]*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"<input[^>]*id=\"(new_)?spawn_rate\"[^>]*>", flags=re.I)
+        re_disabled_spawn_rate = re.compile(
+            r"<input[^>]*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):

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,