Allow setting run time from the web UI / http api

locust/test/test_web.py

@@ -741,6 +741,54 @@ 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": "5s"},
+        )
+        self.assertEqual(200, response.status_code)
+        self.assertEqual("https://localhost", response.json()["host"])
+        self.assertEqual(self.environment.host, "https://localhost")
+        self.assertEqual(5, response.json()["run_time"])
+        # wait for test to run
+        gevent.sleep(7)
+        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_host_value_from_user_class(self):
         class MyUser(User):
             host = "http://example.com"

locust/web.py

@@ -23,6 +23,7 @@
 from .user.inspectuser import get_ratio
 from .util.cache import memoize
 from .util.rounding import proper_round
+from .util.timespan import parse_timespan
 from .html import get_html_report
 from flask_cors import CORS
 
@@ -187,6 +188,7 @@ def swarm() -> Response:
                 self._stop_runners()
 
             parsed_options_dict = vars(environment.parsed_options) if environment.parsed_options else {}
+            run_time = None
             for key, value in request.form.items():
                 if key == "user_count":  # if we just renamed this field to "users" we wouldn't need this
                     user_count = int(value)
@@ -198,6 +200,13 @@ def swarm() -> Response:
                 elif key == "user_classes":
                     # Set environment.parsed_options.user_classes to the selected user_classes
                     parsed_options_dict[key] = request.form.getlist("user_classes")
+                elif key == "run_time":
+                    try:
+                        run_time = parse_timespan(value)
+                    except ValueError:
+                        err_msg = "Valid run_time formats are : 20, 20s, 3m, 2h, 1h20m, 3h30m10s, etc."
+                        logger.error(err_msg)
+                        return jsonify({"success": False, "message": err_msg, "host": environment.host})
                 elif key in parsed_options_dict:
                     # update the value in environment.parsed_options, but dont change the type.
                     # This won't work for parameters that are None
@@ -221,6 +230,9 @@ def swarm() -> Response:
                     "message": "Swarming started",
                     "host": environment.host,
                 }
+                if run_time:
+                    gevent.spawn_later(run_time, self._stop_runners).link_exception(greenlet_exception_handler)
+                    response_data["run_time"] = run_time
 
                 if self.userclass_picker_is_active:
                     response_data["user_classes"] = sorted(user_classes.keys())