diff --git a/locust/argument_parser.py b/locust/argument_parser.py index 2544e736c1..61741b7824 100644 --- a/locust/argument_parser.py +++ b/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", ) diff --git a/locust/main.py b/locust/main.py index f5819d2704..e053680bbe 100644 --- a/locust/main.py +++ b/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), diff --git a/locust/test/test_main.py b/locust/test/test_main.py index cbc6410969..3bdab6099a 100644 --- a/locust/test/test_main.py +++ b/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)