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

…nd users parameters

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

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"
+            "In order to use --run-time, --users or --spawn-rate with a custom LoadShape, you need to define use_common_options=True."
         )
+        logger.warning("See: https://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()

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)

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()
 

locust/templates/index.html

@@ -65,15 +65,15 @@ <h2>Start new load test</h2>
                         {% for user in available_user_classes %}
                             <option value="{{user}}" selected>{{user}}</option>
                         {% endfor %}
-                        </select><br>
+                        </select><br>range
                         <label for="available_shape_classes">ShapeClass </label>
                         <select name="shape_class" id="shape-classes" class="shapeClass">
                         {% for shape_class in available_shape_classes %}
                             <option value="{{shape_class}}">{{shape_class}}</option>
                         {% endfor %}
                         </select><br>
                     {% endif %}
-                    {% if is_shape %}
+                    {% if hide_common_options %}
                         <label for="user_count">Number of users <span style="color:#8a8a8a;">(peak concurrency)</span></label>
                         <input type="text" name="user_count" id="user_count" class="val_disabled" value="-" disabled="disabled" title="Disabled for tests using LoadTestShape class"/><br>
                         <input type="hidden" name="user_count" id="user_count" value="0"/><br>
@@ -144,7 +144,7 @@ <h2>Start new load test</h2>
             <div class="padder">
                 <h2>Edit running load test</h2>
                 <form action="./swarm" method="POST" id="edit_form">
-                    {% if is_shape %}
+                    {% if hide_common_options %}
                         <label for="new_user_count">Number of users (peak concurrency)</label>
                         <input type="text" name="user_count" id="new_user_count" class="val_disabled"  value="-" disabled="disabled" title="Disabled for tests using LoadTestShape class"/><br>
                         <label for="spawn_rate">Spawn rate <span style="color:#8a8a8a;">(users added/stopped per second)</span></label>
@@ -155,7 +155,7 @@ <h2>Edit running load test</h2>
                         <label for="spawn_rate">Spawn rate <span style="color:#8a8a8a;">(users added/stopped per second)</span></label>
                         <input type="text" name="spawn_rate" id="new_spawn_rate" class="val"  value="{{ spawn_rate or "1" }}" onfocus="this.select()"/><br>
                     {% endif %}
-                    {% if is_shape %}
+                    {% if hide_common_options %}
                         <button type="submit" disabled>Start swarming</button>
                     {% else %}
                         <button type="submit">Start swarming</button>

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,98 @@ 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(
+                "In order to use --run-time, --users or --spawn-rate with a custom LoadShape, "
+                "you need to define use_common_options=True.",
+                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(
+                "In order to use --run-time, --users or --spawn-rate with a custom LoadShape, "
+                "you need to define use_common_options=True.",
+                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(
+                "In order to use --run-time, --users or --spawn-rate with a custom LoadShape, "
+                "you need to define use_common_options=True.",
+                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)

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_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"<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,
+            "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,