Merge pull request #1922 from locustio/add-expect-workers-max-wait-pa…

…rameter

Add --expect-workers-max-wait parameter

locust/argument_parser.py

@@ -289,9 +289,17 @@ def setup_parser_arguments(parser):
         "--expect-workers",
         type=int,
         default=1,
-        help="How many workers master should expect to connect before starting the test (only when --headless used).",
+        help="How many workers master should expect to connect before starting the test (only when --headless/autostart is used).",
         env_var="LOCUST_EXPECT_WORKERS",
     )
+    master_group.add_argument(
+        "--expect-workers-max-wait",
+        type=int,
+        default=0,
+        help="How long should the master wait for workers to connect before giving up. Defaults to wait forever",
+        env_var="LOCUST_EXPECT_WORKERS_MAX_WAIT",
+    )
+
     master_group.add_argument(
         "--expect-slaves",
         action="store_true",
@@ -300,12 +308,8 @@ def setup_parser_arguments(parser):
 
     worker_group = parser.add_argument_group(
         "Worker options",
-        textwrap.dedent(
-            """
-            Options for running a Locust Worker node when running Locust distributed.
-            Only the LOCUSTFILE (-f option) need to be specified when starting a Worker, since other options such as -u, -r, -t are specified on the Master node.
-        """
-        ),
+        """Options for running a Locust Worker node when running Locust distributed.
+Only the LOCUSTFILE (-f option) needs to be specified when starting a Worker, since other options such as -u, -r, -t are specified on the Master node.""",
     )
     # if locust should be run in distributed mode as worker
     worker_group.add_argument(
@@ -415,7 +419,7 @@ def setup_parser_arguments(parser):
     )
     log_group.add_argument(
         "--logfile",
-        help="Path to log file. If not set, log will go to stdout/stderr",
+        help="Path to log file. If not set, log will go to stderr",
         env_var="LOCUST_LOGFILE",
     )
 

locust/main.py

@@ -346,7 +346,11 @@ def spawn_run_time_quit_greenlet():
     def start_automatic_run():
         if options.master:
             # wait for worker nodes to connect
+            start_time = time.monotonic()
             while len(runner.clients.ready) < options.expect_workers:
+                if options.expect_workers_max_wait and options.expect_workers_max_wait < time.monotonic() - start_time:
+                    logger.error("Gave up waiting for workers to connect.")
+                    sys.exit(1)
                 logging.info(
                     "Waiting for workers to be ready, %s of %s connected",
                     len(runner.clients.ready),

locust/test/test_main.py

@@ -667,3 +667,26 @@ def test_html_report_option(self):
         self.assertIn("charts-container", html_report_content)
 
         self.assertNotIn("Download the Report", html_report_content, "Download report link found in HTML content")
+
+    def test_expect_workers(self):
+        with mock_locustfile() as mocked:
+            proc = subprocess.Popen(
+                [
+                    "locust",
+                    "-f",
+                    mocked.file_path,
+                    "--headless",
+                    "--master",
+                    "--expect-workers",
+                    "2",
+                    "--expect-workers-max-wait",
+                    "1",
+                ],
+                stdout=PIPE,
+                stderr=PIPE,
+            )
+            _, stderr = proc.communicate()
+            stderr = stderr.decode("utf-8")
+            self.assertIn("Waiting for workers to be ready, 0 of 2 connected", stderr)
+            self.assertIn("Gave up waiting for workers to connect", stderr)
+            self.assertEqual(1, proc.returncode)