Allow setting run time from the web UI / http api

docs/custom-load-shape.rst

@@ -41,11 +41,21 @@ This functionality is further demonstrated in the `examples on github <https://g
 
 One further method may be helpful for your custom load shapes: `get_current_user_count()`, which returns the total number of active users. This method can be used to prevent advancing to subsequent steps until the desired number of users has been reached. This is especially useful if the initialization process for each user is slow or erratic in how long it takes. If this sounds like your use case, see the `example on github <https://github.com/locustio/locust/tree/master/examples/custom_shape/wait_user_count.py>`_.
 
+Combining Users with different load profiles
+--------------------------------------------
 
-Extend your shape with custom users
------------------------------------
+If you use the Web UI, you can add the :ref:`---class-picker <class-picker>` parameter to select which shape to use. But it often more flexible to have your User definitions in one file and your LoadTestShape in a separate one. For example, if you a high/low load Shape class defined in low_load.py and high_load.py respectively:
 
-Extending the return value of the ``tick()`` with the argument ``user_classes`` makes it possible to pick the users being created for a ``tick()`` specifically.
+.. code-block:: console
+
+    $ locust -f locustfile.py,low_load.py
+
+    $ locust -f locustfile.py,high_load.py
+
+Restricting which user types to spawn in each tick
+--------------------------------------------------
+
+Adding the element ``user_classes`` to the return value gives you more detailed control:
 
 .. code-block:: python
 

docs/quickstart.rst

@@ -123,6 +123,7 @@ Example:
 
 Locust will use ``locustfile1.py``, ``locustfile2.py`` & ``more_files/locustfile3.py``
 
+.. _class-picker:
 
 Running Locust with User class UI picker
 ========================================

locust/clients.py

@@ -1,7 +1,8 @@
+from __future__ import annotations
 import re
 import time
 from contextlib import contextmanager
-from typing import Generator, Optional, Union
+from typing import Generator, Optional
 from urllib.parse import urlparse, urlunparse
 
 import requests
@@ -220,7 +221,7 @@ class ResponseContextManager(LocustResponse):
     :py:meth:`failure <locust.clients.ResponseContextManager.failure>`.
     """
 
-    _manual_result: Optional[Union[bool, Exception]] = None
+    _manual_result: bool | Exception | None = None
     _entered = False
 
     def __init__(self, response, request_event, request_meta):

locust/contrib/fasthttp.py

@@ -7,7 +7,7 @@
 from urllib.parse import urlparse, urlunparse
 from ssl import SSLError
 import time
-from typing import Callable, Optional, Tuple, Dict, Any, Union
+from typing import Callable, Optional, Tuple, Dict, Any
 
 from http.cookiejar import CookieJar
 
@@ -143,17 +143,17 @@ def request(
         self,
         method: str,
         url: str,
-        name: str = None,
-        data: Union[str, dict] = None,
+        name: str | None = None,
+        data: str | dict | None = None,
         catch_response: bool = False,
         stream: bool = False,
-        headers: dict = None,
+        headers: dict | None = None,
         auth=None,
-        json: dict = None,
+        json: dict | None = None,
         allow_redirects=True,
         context: dict = {},
         **kwargs,
-    ) -> Union[ResponseContextManager, FastResponse]:
+    ) -> ResponseContextManager | FastResponse:
         """
         Send and HTTP request
         Returns :py:class:`locust.contrib.fasthttp.FastResponse` object.

locust/env.py

@@ -6,7 +6,6 @@
     Type,
     TypeVar,
     Optional,
-    Union,
 )
 
 from configargparse import Namespace
@@ -239,7 +238,7 @@ def assign_equal_weights(self) -> None:
         """
         for u in self.user_classes:
             u.weight = 1
-            user_tasks: List[Union[TaskSet, Callable]] = []
+            user_tasks: List[TaskSet | Callable] = []
             tasks_frontier = u.tasks
             while len(tasks_frontier) != 0:
                 t = tasks_frontier.pop()

locust/runners.py

@@ -28,7 +28,6 @@
     Type,
     Any,
     cast,
-    Union,
 )
 from uuid import uuid4
 
@@ -118,7 +117,7 @@ def __init__(self, environment: "Environment") -> None:
         self.state = STATE_INIT
         self.spawning_greenlet: Optional[gevent.Greenlet] = None
         self.shape_greenlet: Optional[gevent.Greenlet] = None
-        self.shape_last_tick: Union[Tuple[int, float], Tuple[int, float, Optional[List[Type[User]]]], None] = None
+        self.shape_last_tick: Tuple[int, float] | Tuple[int, float, Optional[List[Type[User]]]] | None = None
         self.current_cpu_usage: int = 0
         self.cpu_warning_emitted: bool = False
         self.worker_cpu_warning_emitted: bool = False
@@ -352,9 +351,7 @@ def start_shape(self) -> None:
     def shape_worker(self) -> None:
         logger.info("Shape worker starting")
         while self.state == STATE_INIT or self.state == STATE_SPAWNING or self.state == STATE_RUNNING:
-            current_tick: Union[Tuple[int, float], Tuple[int, float, Optional[List[Type[User]]]], None] = (
-                self.environment.shape_class.tick() if self.environment.shape_class is not None else None
-            )
+            current_tick = self.environment.shape_class.tick() if self.environment.shape_class is not None else None
             if current_tick is None:
                 logger.info("Shape test stopping")
                 if self.environment.parsed_options and self.environment.parsed_options.headless:

locust/shape.py

@@ -1,5 +1,6 @@
+from __future__ import annotations
 import time
-from typing import Optional, Tuple, List, Type, Union
+from typing import Optional, Tuple, List, Type
 
 from . import User
 from .runners import Runner
@@ -35,7 +36,7 @@ def get_current_user_count(self):
         """
         return self.runner.user_count
 
-    def tick(self) -> Union[Tuple[int, float], Tuple[int, float, Optional[List[Type[User]]]], None]:
+    def tick(self) -> Tuple[int, float] | Tuple[int, float, Optional[List[Type[User]]]] | None:
         """
         Returns a tuple with 2 elements to control the running load test:
 

locust/stats.py

@@ -1,3 +1,4 @@
+from __future__ import annotations
 from abc import abstractmethod
 import datetime
 import hashlib
@@ -19,7 +20,6 @@
     NoReturn,
     Tuple,
     List,
-    Union,
     Optional,
     OrderedDict as OrderedDictType,
     Callable,
@@ -57,7 +57,7 @@
 
 class CSVWriter(Protocol):
     @abstractmethod
-    def writerow(self, columns: Iterable[Union[str, int, float]]) -> None:
+    def writerow(self, columns: Iterable[str | int | float]) -> None:
         ...
 
 
@@ -228,7 +228,7 @@ def log_request(self, method: str, name: str, response_time: int, content_length
         self.total.log(response_time, content_length)
         self.get(name, method).log(response_time, content_length)
 
-    def log_error(self, method: str, name: str, error: Optional[Union[Exception, str]]) -> None:
+    def log_error(self, method: str, name: str, error: Exception | str | None) -> None:
         self.total.log_error(error)
         self.get(name, method).log_error(error)
 
@@ -404,7 +404,7 @@ def _log_response_time(self, response_time: int) -> None:
         self.response_times.setdefault(rounded_response_time, 0)
         self.response_times[rounded_response_time] += 1
 
-    def log_error(self, error: Optional[Union[Exception, str]]) -> None:
+    def log_error(self, error: Exception | str | None) -> None:
         self.num_failures += 1
         t = int(time.time())
         self.num_fail_per_sec[t] = self.num_fail_per_sec.setdefault(t, 0) + 1
@@ -449,7 +449,7 @@ def current_rps(self) -> float:
             return 0
         slice_start_time = max(int(self.stats.last_request_timestamp) - 12, int(self.stats.start_time or 0))
 
-        reqs: List[Union[int, float]] = [
+        reqs: List[int | float] = [
             self.num_reqs_per_sec.get(t, 0) for t in range(slice_start_time, int(self.stats.last_request_timestamp) - 2)
         ]
         return avg(reqs)
@@ -679,14 +679,14 @@ def _cache_response_times(self, t: int) -> None:
 
 
 class StatsError:
-    def __init__(self, method: str, name: str, error: Optional[Union[Exception, str]], occurrences: int = 0):
+    def __init__(self, method: str, name: str, error: Exception | str | None, occurrences: int = 0):
         self.method = method
         self.name = name
         self.error = error
         self.occurrences = occurrences
 
     @classmethod
-    def parse_error(cls, error: Optional[Union[Exception, str]]) -> str:
+    def parse_error(cls, error: Exception | str | None) -> str:
         string_error = repr(error)
         target = "object at 0x"
         target_index = string_error.find(target)
@@ -700,7 +700,7 @@ def parse_error(cls, error: Optional[Union[Exception, str]]) -> str:
         return string_error.replace(hex_address, "0x....")
 
     @classmethod
-    def create_key(cls, method: str, name: str, error: Optional[Union[Exception, str]]) -> str:
+    def create_key(cls, method: str, name: str, error: Exception | str | None) -> str:
         key = f"{method}.{name}.{StatsError.parse_error(error)!r}"
         return hashlib.sha256(key.encode("utf-8")).hexdigest()
 
@@ -738,7 +738,7 @@ def unserialize(cls, data: StatsErrorDict) -> "StatsError":
         return cls(data["method"], data["name"], data["error"], data["occurrences"])
 
 
-def avg(values: List[Union[float, int]]) -> float:
+def avg(values: List[float | int]) -> float:
     return sum(values, 0.0) / max(len(values), 1)
 
 
@@ -906,7 +906,7 @@ def __init__(self, environment: "Environment", percentiles_to_report: List[float
             "Nodes",
         ]
 
-    def _percentile_fields(self, stats_entry: StatsEntry, use_current: bool = False) -> Union[List[str], List[int]]:
+    def _percentile_fields(self, stats_entry: StatsEntry, use_current: bool = False) -> List[str] | List[int]:
         if not stats_entry.num_requests:
             return self.percentiles_na
         elif use_current:

locust/templates/index.html

@@ -93,6 +93,14 @@ <h2>Start new load test</h2>
                         {% endif %}
                     </label>
                     <input type="text" name="host" id="host" class="val" autocapitalize="off" autocorrect="off" value="{{ host  or "" }}" onfocus="this.select()"/><br>
+                    <a href="#" onclick="$('.advancedOptions').toggle();">Advanced options</a>
+                    <div class="advancedOptions" style="display:none">
+                        <label for="run_time">
+                            Run time
+                            <span style="color:#8a8a8a;">(e.g. 20, 20s, 3m, 2h, 1h20m, 3h30m10s, etc.)</span>
+                        </label>
+                        <input type="text" name="run_time" id="run_time" class="val" value="{{ run_time  or "" }}" onfocus="this.select()"/><br>
+                    </div>
                     {% if extra_options %}<label>Custom parameters:</label>{% endif %}
                     {% for key, value in extra_options.items() %}
                         {% if not ((value is none) or (value is boolean)) %}

locust/test/test_web.py

@@ -741,6 +741,81 @@ def my_task(self):
         self.assertEqual(None, response.json()["host"])
         self.assertEqual(self.environment.host, None)
 
+    def test_swarm_run_time(self):
+        class MyUser(User):
+            wait_time = constant(1)
+
+            @task(1)
+            def my_task(self):
+                pass
+
+        self.environment.user_classes = [MyUser]
+        self.environment.web_ui.parsed_options = parse_options()
+        response = requests.post(
+            "http://127.0.0.1:%i/swarm" % self.web_port,
+            data={"user_count": 5, "spawn_rate": 5, "host": "https://localhost", "run_time": "1s"},
+        )
+        self.assertEqual(200, response.status_code)
+        self.assertEqual("https://localhost", response.json()["host"])
+        self.assertEqual(self.environment.host, "https://localhost")
+        self.assertEqual(1, response.json()["run_time"])
+        # wait for test to run
+        gevent.sleep(3)
+        response = requests.get("http://127.0.0.1:%i/stats/requests" % self.web_port)
+        self.assertEqual("stopped", response.json()["state"])
+
+    def test_swarm_run_time_invalid_input(self):
+        class MyUser(User):
+            wait_time = constant(1)
+
+            @task(1)
+            def my_task(self):
+                pass
+
+        self.environment.user_classes = [MyUser]
+        self.environment.web_ui.parsed_options = parse_options()
+        response = requests.post(
+            "http://127.0.0.1:%i/swarm" % self.web_port,
+            data={"user_count": 5, "spawn_rate": 5, "host": "https://localhost", "run_time": "bad"},
+        )
+        self.assertEqual(200, response.status_code)
+        self.assertEqual(False, response.json()["success"])
+        self.assertEqual(self.environment.host, "https://localhost")
+        self.assertEqual(
+            "Valid run_time formats are : 20, 20s, 3m, 2h, 1h20m, 3h30m10s, etc.", response.json()["message"]
+        )
+        # verify test was not started
+        response = requests.get("http://127.0.0.1:%i/stats/requests" % self.web_port)
+        self.assertEqual("ready", response.json()["state"])
+        requests.get("http://127.0.0.1:%i/stats/reset" % self.web_port)
+
+    def test_swarm_run_time_empty_input(self):
+        class MyUser(User):
+            wait_time = constant(1)
+
+            @task(1)
+            def my_task(self):
+                pass
+
+        self.environment.user_classes = [MyUser]
+        self.environment.web_ui.parsed_options = parse_options()
+        response = requests.post(
+            "http://127.0.0.1:%i/swarm" % self.web_port,
+            data={"user_count": 5, "spawn_rate": 5, "host": "https://localhost", "run_time": ""},
+        )
+
+        self.assertEqual(200, response.status_code)
+        self.assertEqual("https://localhost", response.json()["host"])
+        self.assertEqual(self.environment.host, "https://localhost")
+
+        # verify test is running
+        gevent.sleep(1)
+        response = requests.get("http://127.0.0.1:%i/stats/requests" % self.web_port)
+        self.assertEqual("running", response.json()["state"])
+
+        # stop
+        response = requests.get("http://127.0.0.1:%i/stop" % self.web_port)
+
     def test_host_value_from_user_class(self):
         class MyUser(User):
             host = "http://example.com"

locust/user/task.py

@@ -1,3 +1,4 @@
+from __future__ import annotations
 import logging
 import random
 import traceback
@@ -6,7 +7,6 @@
     TYPE_CHECKING,
     Callable,
     List,
-    Union,
     TypeVar,
     Optional,
     Type,
@@ -51,7 +51,7 @@ def task(weight: int) -> Callable[[TaskT], TaskT]:
     ...
 
 
-def task(weight: Union[TaskT, int] = 1) -> Union[TaskT, Callable[[TaskT], TaskT]]:
+def task(weight: TaskT | int = 1) -> TaskT | Callable[[TaskT], TaskT]:
     """
     Used as a convenience decorator to be able to declare tasks for a User or a TaskSet
     inline in the class. Example::
@@ -242,7 +242,7 @@ class TaskSet(metaclass=TaskSetMeta):
     will then continue in the first TaskSet).
     """
 
-    tasks: List[Union["TaskSet", Callable]] = []
+    tasks: List[TaskSet | Callable] = []
     """
     Collection of python callables and/or TaskSet classes that the User(s) will run.
 

locust/user/users.py

@@ -1,4 +1,5 @@
-from typing import Callable, Dict, List, Optional, Union
+from __future__ import annotations
+from typing import Callable, Dict, List, Optional
 
 from gevent import GreenletExit, greenlet
 from gevent.pool import Group
@@ -83,7 +84,7 @@ class MyUser(User):
     Method that returns the time between the execution of locust tasks in milliseconds
     """
 
-    tasks: List[Union[TaskSet, Callable]] = []
+    tasks: List[TaskSet | Callable] = []
     """
     Collection of python callables and/or TaskSet classes that the Locust user(s) will run.