Add --autostart and --autoquit parameters, fixes #831

locust/argument_parser.py

@@ -141,14 +141,14 @@ def setup_parser_arguments(parser):
         "--users",
         type=int,
         dest="num_users",
-        help="Number of concurrent Locust users. Primarily used together with --headless. Can be changed during a test by inputs w, W(spawn 1, 10 users) and s, S(stop 1, 10 users)",
+        help="Peak number of concurrent Locust users. Primarily used together with --headless or --autostart. Can be changed during a test by keyboard inputs w, W (spawn 1, 10 users) and s, S (stop 1, 10 users)",
         env_var="LOCUST_USERS",
     )
     parser.add_argument(
         "-r",
         "--spawn-rate",
         type=float,
-        help="The rate per second in which users are spawned. Primarily used together with --headless",
+        help="Rate to spawn users at (users per second). Primarily used together with --headless or --autostart",
         env_var="LOCUST_SPAWN_RATE",
     )
     parser.add_argument(
@@ -161,7 +161,7 @@ def setup_parser_arguments(parser):
     parser.add_argument(
         "-t",
         "--run-time",
-        help="Stop after the specified amount of time, e.g. (300s, 20m, 3h, 1h30m, etc.). Only used together with --headless. Defaults to run forever.",
+        help="Stop after the specified amount of time, e.g. (300s, 20m, 3h, 1h30m, etc.). Only used together with --headless or --autostart. Defaults to run forever.",
         env_var="LOCUST_RUN_TIME",
     )
     parser.add_argument(
@@ -190,9 +190,22 @@ def setup_parser_arguments(parser):
     web_ui_group.add_argument(
         "--headless",
         action="store_true",
-        help="Disable the web interface, and instead start the load test immediately. Requires -u and -t to be specified.",
+        help="Disable the web interface, and start the test immediately. Use -u and -t to control user count and run time",
         env_var="LOCUST_HEADLESS",
     )
+    web_ui_group.add_argument(
+        "--autostart",
+        action="store_true",
+        help="Starts the test immediately (without disabling the web UI). Use -u and -t to control user count and run time",
+        env_var="LOCUST_AUTOSTART",
+    )
+    web_ui_group.add_argument(
+        "--autoquit",
+        type=int,
+        default=-1,
+        help="Quits Locust entirely, X seconds after the run is finished. Only used together with --autostart. The default is to keep Locust running until you shut it down using CTRL+C",
+        env_var="LOCUST_AUTOQUIT",
+    )
     # Override --headless parameter (useful because you cant disable a store_true-parameter like headless once it has been set in a config file)
     web_ui_group.add_argument(
         "--headful",

locust/input_events.py

@@ -93,7 +93,7 @@ def input_listener_func():
         try:
             poller = get_poller()
         except InitError as e:
-            logging.info(e)
+            logging.debug(e)
             return
 
         try:

locust/main.py

@@ -134,6 +134,10 @@ def main():
         sys.stderr.write("The --slave/--expect-slaves parameters have been renamed --worker/--expect-workers\n")
         sys.exit(1)
 
+    if options.autoquit != -1 and not options.autostart:
+        sys.stderr.write("--autoquit is only meaningful in combination with --autostart\n")
+        sys.exit(1)
+
     if options.step_time or options.step_load or options.step_users or options.step_clients:
         sys.stderr.write(
             "The step load feature was removed in Locust 1.3. You can achieve similar results using a LoadTestShape class. See https://docs.locust.io/en/stable/custom-load-shape.html\n"
@@ -247,9 +251,6 @@ def main():
     main_greenlet = runner.greenlet
 
     if options.run_time:
-        if not options.headless:
-            logger.error("The --run-time argument can only be used together with --headless")
-            sys.exit(1)
         if options.worker:
             logger.error("--run-time should be specified on the master node, and not on worker nodes")
             sys.exit(1)
@@ -312,6 +313,13 @@ def assign_equal_weights(environment, **kwargs):
         web_ui.start()
         main_greenlet = web_ui.greenlet
 
+    def spawn_run_time_quit_greenlet():
+        def timelimit_quit():
+            logger.info("--run-time limit reached. Stopping Locust")
+            runner.quit()
+
+        gevent.spawn_later(options.run_time, timelimit_quit).link_exception(greenlet_exception_handler)
+
     headless_master_greenlet = None
     if options.headless:
         # headless mode
@@ -336,25 +344,19 @@ def assign_equal_weights(environment, **kwargs):
 
             # 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)
                 environment.runner.start_shape()
             else:
                 headless_master_greenlet = gevent.spawn(runner.start, options.num_users, options.spawn_rate)
                 headless_master_greenlet.link_exception(greenlet_exception_handler)
 
-    def spawn_run_time_limit_greenlet():
-        def timelimit_stop():
-            logger.info("Time limit reached. Stopping Locust.")
-            runner.quit()
-
-        gevent.spawn_later(options.run_time, timelimit_stop).link_exception(greenlet_exception_handler)
-
-    if options.run_time:
-        logger.info("Run time limit set to %s seconds" % options.run_time)
-        spawn_run_time_limit_greenlet()
-    elif options.headless and not options.worker and not environment.shape_class:
-        logger.info("No run time limit set, use CTRL+C to interrupt.")
-    else:
-        pass  # dont log anything - not having a time limit is normal when not running headless
+        if options.run_time:
+            logger.info("Run time limit set to %s seconds" % options.run_time)
+            spawn_run_time_quit_greenlet()
+        elif not options.worker and not environment.shape_class:
+            logger.info("No run time limit set, use CTRL+C to interrupt")
 
     input_listener_greenlet = None
     if not options.worker:
@@ -434,13 +436,56 @@ def sig_term_handler():
 
     gevent.signal_handler(signal.SIGTERM, sig_term_handler)
 
+    def autoquit():
+        if options.autoquit != -1:
+            logger.debug("Autoquit time limit set to %s seconds" % options.autoquit)
+            time.sleep(options.autoquit)
+            logger.info("--autoquit time reached, shutting down")
+            runner.quit()
+            web_ui.stop()
+        else:
+            logger.info("--autoquit not specified, leaving web ui running indefinitely")
+
     try:
         logger.info("Starting Locust %s" % version)
+        if options.autostart:
+            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)
+                environment.runner.start_shape()
+                autoquit()
+            else:
+                if not options.worker:
+                    if options.num_users is None:
+                        options.num_users = 1
+                    if options.spawn_rate is None:
+                        options.spawn_rate = 1
+
+                autostart_master_greenlet = gevent.spawn(runner.start, options.num_users, options.spawn_rate)
+                autostart_master_greenlet.link_exception(greenlet_exception_handler)
+
+                def timelimit_stop():
+                    logger.info("--run-time limit reached, stopping test")
+                    runner.stop()
+                    autoquit()
+
+                if options.run_time:
+                    logger.info("Run time limit set to %s seconds" % options.run_time)
+                    gevent.spawn_later(options.run_time, timelimit_stop).link_exception(greenlet_exception_handler)
+                else:
+                    logger.info("No run time limit set, use CTRL+C to interrupt")
+
         main_greenlet.join()
         if options.html_file:
             html_report = get_html_report(environment, show_download_link=False)
             with open(options.html_file, "w", encoding="utf-8") as file:
                 file.write(html_report)
-        shutdown()
     except KeyboardInterrupt:
-        shutdown()
+        pass
+    except Exception:
+        if input_listener_greenlet is not None:
+            input_listener_greenlet.kill(block=False)
+            time.sleep(0)
+        raise
+    shutdown()

locust/test/test_log.py

@@ -75,7 +75,7 @@ def my_task(self):
             output,
         )
         self.assertIn(
-            "%s/INFO/locust.main: Time limit reached. Stopping Locust." % socket.gethostname(),
+            "%s/INFO/locust.main: --run-time limit reached. Stopping Locust" % socket.gethostname(),
             output,
         )
         self.assertIn(
@@ -187,7 +187,7 @@ def my_task(self):
             log_content,
         )
         self.assertIn(
-            "%s/INFO/locust.main: Time limit reached. Stopping Locust." % socket.gethostname(),
+            "%s/INFO/locust.main: --run-time limit reached. Stopping Locust" % socket.gethostname(),
             log_content,
         )
         self.assertIn(

locust/test/test_main.py

@@ -18,6 +18,19 @@
 from .testcases import LocustTestCase
 from .util import temporary_file, get_free_tcp_port
 
+SIMPLE_LOCUST_FILE = textwrap.dedent(
+    """
+    from locust import HttpUser, task
+    import time
+    class UserSubclass(HttpUser):
+        host = "https://www.test.com"
+        @task
+        def t(self):
+            self.client.get("/")
+            time.sleep(1)
+    """
+)
+
 
 class TestLoadLocustfile(LocustTestCase):
     def test_is_user_class(self):
@@ -198,11 +211,11 @@ def my_task(self):
             gevent.sleep(1)
             proc.send_signal(signal.SIGTERM)
             stdout, stderr = proc.communicate()
-            self.assertEqual(42, proc.returncode)
             stderr = stderr.decode("utf-8")
             self.assertIn("Starting web interface at", stderr)
             self.assertIn("Starting Locust", stderr)
             self.assertIn("Shutting down (exit code 42), bye", stderr)
+            self.assertEqual(42, proc.returncode)
 
     def test_webserver(self):
         with temporary_file(
@@ -282,7 +295,7 @@ def tick(self):
         with mock_locustfile(content=content) as mocked:
             output = (
                 subprocess.check_output(
-                    ["locust", "-f", mocked.file_path, "--host", "https://test.com/", "--run-time", "1s", "--headless"],
+                    ["locust", "-f", mocked.file_path, "--host", "https://test.com/", "--headless"],
                     stderr=subprocess.STDOUT,
                     timeout=3,
                 )
@@ -292,6 +305,114 @@ def tick(self):
             self.assertIn("Shape test updating to 10 users at 1.00 spawn rate", output)
             self.assertIn("Cleaning up runner...", output)
 
+    def test_autostart_wo_run_time(self):
+        port = get_free_tcp_port()
+        with mock_locustfile(content=SIMPLE_LOCUST_FILE) as mocked:
+            proc = subprocess.Popen(
+                [
+                    "locust",
+                    "-f",
+                    mocked.file_path,
+                    "--web-port",
+                    str(port),
+                    "--autostart",
+                ],
+                stdout=PIPE,
+                stderr=PIPE,
+            )
+            gevent.sleep(1.8)
+            response = requests.get(f"http://0.0.0.0:{port}/stats/requests")
+            self.assertEqual(200, response.status_code)
+            proc.send_signal(signal.SIGTERM)
+            stdout, stderr = proc.communicate()
+            stderr = stderr.decode("utf-8")
+            self.assertIn("Starting Locust", stderr)
+            self.assertIn("No run time limit set, use CTRL+C to interrupt", stderr)
+            self.assertIn("Shutting down (exit code 0), bye", stderr)
+            self.assertNotIn("Traceback", stderr)
+            # check stats afterwards, because it really isnt as informative as the output itself
+            data = response.json()
+            self.assertEqual(2, len(data["stats"]), data)
+            self.assertEqual("/", data["stats"][0]["name"])
+
+    def test_autostart_w_run_time(self):
+        port = get_free_tcp_port()
+        with mock_locustfile(content=SIMPLE_LOCUST_FILE) as mocked:
+            proc = subprocess.Popen(
+                [
+                    "locust",
+                    "-f",
+                    mocked.file_path,
+                    "--web-port",
+                    str(port),
+                    "-t",
+                    "1",
+                    "--autostart",
+                    "--autoquit",
+                    "2",
+                ],
+                stdout=PIPE,
+                stderr=PIPE,
+            )
+            gevent.sleep(2.9)
+            response = requests.get(f"http://0.0.0.0:{port}/stats/requests")
+            self.assertEqual(200, response.status_code)
+            _, stderr = proc.communicate(timeout=2)
+            stderr = stderr.decode("utf-8")
+            self.assertIn("Starting Locust", stderr)
+            self.assertIn("Run time limit set to 1 seconds", stderr)
+            self.assertIn("Shutting down (exit code 0), bye", stderr)
+            self.assertNotIn("Traceback", stderr)
+            data = response.json()
+            # check stats afterwards, because it really isnt as informative as the output itself
+            self.assertEqual(2, len(data["stats"]), data)
+            self.assertEqual("/", data["stats"][0]["name"])
+
+    def test_autostart_w_load_shape(self):
+        port = get_free_tcp_port()
+        with mock_locustfile(
+            content=SIMPLE_LOCUST_FILE
+            + textwrap.dedent(
+                """
+            from locust import LoadTestShape
+            class LoadTestShape(LoadTestShape):
+                def tick(self):
+                    run_time = self.get_run_time()
+                    if run_time < 2:
+                        return (10, 1)
+
+                    return None
+            """
+            )
+        ) as mocked:
+            proc = subprocess.Popen(
+                [
+                    "locust",
+                    "-f",
+                    mocked.file_path,
+                    "--web-port",
+                    str(port),
+                    "--autostart",
+                    "--autoquit",
+                    "2",
+                ],
+                stdout=PIPE,
+                stderr=PIPE,
+            )
+            gevent.sleep(1.8)
+            response = requests.get(f"http://0.0.0.0:{port}/stats/requests")
+            self.assertEqual(200, response.status_code)
+            _, stderr = proc.communicate(timeout=4)
+            stderr = stderr.decode("utf-8")
+            self.assertIn("Starting Locust", stderr)
+            self.assertIn("Shape test starting", stderr)
+            self.assertIn("Shutting down (exit code 0), bye", stderr)
+            self.assertNotIn("Traceback", stderr)
+            # check stats afterwards, because it really isnt as informative as the output itself
+            data = response.json()
+            self.assertEqual(2, len(data["stats"]), data)
+            self.assertEqual("/", data["stats"][0]["name"])
+
     def test_web_options(self):
         port = get_free_tcp_port()
         if platform.system() == "Darwin":

locust/test/util.py

@@ -34,6 +34,7 @@ def get_free_tcp_port():
     """
     s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
     s.bind(("127.0.0.1", 0))
+    s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
     port = s.getsockname()[1]
     s.close()
     return port