Merge pull request #3009 from locustio/treat-exceptions-in-init-event…

…-handler-as-fatal

Treat exceptions in init event handler as fatal

locust/main.py

@@ -544,6 +544,9 @@ def assign_equal_weights(environment, **kwargs):
     # Fire locust init event which can be used by end-users' code to run setup code that
     # need access to the Environment, Runner or WebUI.
     environment.events.init.fire(environment=environment, runner=runner, web_ui=web_ui)
+    if log.unhandled_greenlet_exception:
+        # treat exceptions in init handlers as fatal. They are already logged so no need to log anything more.
+        sys.exit(1)
 
     if web_ui:
         web_ui.start()

locust/test/test_main.py

@@ -1442,6 +1442,39 @@ def my_task(self):
             # ensure stats printer printed at least one report before shutting down and that there was a final report printed as well
             self.assertRegex(stderr, r".*Aggregated[\S\s]*Shutting down[\S\s]*Aggregated.*")
 
+    def test_exception_in_init_event(self):
+        with mock_locustfile(
+            content=textwrap.dedent(
+                """
+            from locust import User, task, constant, events
+            @events.init.add_listener
+            def _(*args, **kw):
+                raise Exception("something went wrong")
+
+            class TestUser(User):
+                wait_time = constant(10)
+                @task
+                def my_task(self):
+                    pass
+                    """
+            )
+        ) as mocked:
+            proc = subprocess.Popen(
+                [
+                    "locust",
+                    "-f",
+                    mocked.file_path,
+                    "--host",
+                    "https://test.com/",
+                ],
+                stdout=PIPE,
+                stderr=PIPE,
+                text=True,
+            )
+            stdout, stderr = proc.communicate(timeout=3)
+            self.assertIn("something went wrong", stderr)
+            self.assertEqual(1, proc.returncode)
+
 
 class DistributedIntegrationTests(ProcessIntegrationTest):
     failed_port_check = False