Merge pull request #2202 from ajt89/run-time-input-web

Allow setting run time from the web UI / http api

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/web.py

@@ -24,6 +24,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
 
@@ -188,6 +189,7 @@ def swarm() -> Response:
                 self._reset_users_dispatcher()
 
             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)
@@ -199,6 +201,15 @@ 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":
+                    if not value:
+                        continue
+                    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
@@ -222,6 +233,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())
@@ -567,5 +581,8 @@ def _update_user_classes(self, user_classes):
         self.environment._remove_user_classes_with_weight_zero()
         self.environment._validate_user_class_name_uniqueness()
 
+    def _stop_runners(self):
+        self.environment.runner.stop()
+
     def _reset_users_dispatcher(self):
         self.environment.runner._users_dispatcher = None