Merge pull request #1396 from locustio/custom-exit-code

Ability to control the Locust process' exit code

docs/changelog.rst

@@ -4,6 +4,12 @@ Changelog Highlights
 
 For full details of the Locust changelog, please see https://github.com/locustio/locust/blob/master/CHANGELOG.md
 
+In development
+==============
+
+* Ability to control the exit code of the Locust process by setting :py:attr:`Environment.process_exit_code <locust.env.Environment.process_exit_code>`
+
+
 1.0.2
 =====
 

docs/running-locust-without-web-ui.rst

@@ -26,6 +26,7 @@ If you want to specify the run time for a test, you can do that with ``--run-tim
 
 Locust will shutdown once the time is up.
 
+
 Allow tasks to finish their iteration on shutdown
 -------------------------------------------------
 
@@ -37,6 +38,7 @@ By default, locust will stop your tasks immediately. If you want to allow your t
 
 .. _running-locust-distributed-without-web-ui:
 
+
 Running Locust distributed without the web UI
 ---------------------------------------------
 
@@ -45,3 +47,38 @@ you should specify the ``--expect-workers`` option when starting the master node
 the number of worker nodes that are expected to connect. It will then wait until that many worker
 nodes have connected before starting the test.
 
+
+Controlling the exit code of the Locust process
+-----------------------------------------------
+
+When running Locust in a CI environment, you might want to control the exit code of the Locust 
+process. You can do that in your test scripts by setting the 
+:py:attr:`process_exit_code <locust.env.Environment.process_exit_code>` of the 
+:py:class:`Environment <locust.env.Environment>` instance.
+
+Below is an example that'll set the exit code to non zero if any of the following conditions are met:
+
+* More than 1% of the requests failed
+* The average response time is longer than 200 ms
+* The 95th percentile for response time is larger than 800 ms
+
+.. code-block:: python
+
+    import logging
+    from locust import events
+    
+    @events.quitting.add_listener
+    def _(environment, **kw):
+        if environment.stats.total.fail_ratio > 0.01:
+            logging.error("Test failed due to failure ratio > 1%")
+            environment.process_exit_code = 1
+        elif environment.stats.total.avg_response_time > 200:
+            logging.error("Test failed due to average response time ratio > 200 ms")
+            environment.process_exit_code = 1
+        elif environment.stats.total.get_response_time_percentile(0.95) > 800:
+            logging.error("Test failed due to 95th percentil response time > 800 ms")
+            environment.process_exit_code = 1
+        else:
+            environment.process_exit_code = 0
+
+(this code could go into the locustfile.py or in any other file that is imported in the locustfile)

locust/env.py

@@ -51,6 +51,11 @@ class Environment:
     If True exceptions that happen within running users will be catched (and reported in UI/console).
     If False, exeptions will be raised.
     """
+
+    process_exit_code: int = None
+    """
+    If set it'll be the exit code of the Locust process
+    """
 
     parsed_options = None
     """Optional reference to the parsed command line options (used to pre-populate fields in Web UI)"""

locust/event.py

@@ -111,6 +111,10 @@ class Events:
     quitting = EventHook
     """
     Fired when the locust process is exiting
+    
+    Event arguments:
+    
+    :param environment: Environment instance
     """
 
     init = EventHook

locust/main.py

@@ -297,18 +297,30 @@ def timelimit_stop():
         gevent.spawn(stats_writer, environment, options.csv_prefix, full_history=options.stats_history_enabled).link_exception(greenlet_exception_handler)
 
 
-    def shutdown(code=0):
+    def shutdown():
         """
         Shut down locust by firing quitting event, printing/writing stats and exiting
         """
+        logger.info("Running teardowns...")
+        environment.events.quitting.fire(environment=environment, reverse=True)
+
+        # determine the process exit code
+        if log.unhandled_greenlet_exception:
+            code = 2
+        elif environment.process_exit_code is not None:
+            code = environment.process_exit_code
+        elif len(runner.errors) or len(runner.exceptions):
+            code = options.exit_code_on_error
+        else:
+            code = 0
+
         logger.info("Shutting down (exit code %s), bye." % code)
         if stats_printer_greenlet is not None:
             stats_printer_greenlet.kill(block=False)
         logger.info("Cleaning up runner...")
         if runner is not None:
             runner.quit()
-        logger.info("Running teardowns...")
-        environment.events.quitting.fire(reverse=True)
+
         print_stats(runner.stats, current=False)
         print_percentile_stats(runner.stats)
         if options.csv_prefix:
@@ -319,17 +331,12 @@ def shutdown(code=0):
     # install SIGTERM handler
     def sig_term_handler():
         logger.info("Got SIGTERM signal")
-        shutdown(0)
+        shutdown()
     gevent.signal_handler(signal.SIGTERM, sig_term_handler)
 
     try:
         logger.info("Starting Locust %s" % version)
         main_greenlet.join()
-        code = 0
-        if log.unhandled_greenlet_exception:
-            code = 2
-        elif len(runner.errors) or len(runner.exceptions):
-            code = options.exit_code_on_error
-        shutdown(code=code)
+        shutdown()
     except KeyboardInterrupt as e:
-        shutdown(0)
+        shutdown()

locust/runners.py

@@ -413,7 +413,7 @@ def on_worker_report(client_id, data):
         self.environment.events.worker_report.add_listener(on_worker_report)
 
         # register listener that sends quit message to worker nodes
-        def on_quitting():
+        def on_quitting(environment, **kw):
             self.quit()
         self.environment.events.quitting.add_listener(on_quitting)
 
@@ -615,7 +615,7 @@ def on_report_to_master(client_id, data):
         self.environment.events.report_to_master.add_listener(on_report_to_master)
 
         # register listener that sends quit message to master
-        def on_quitting():
+        def on_quitting(environment, **kw):
             self.client.send(Message("quit", None, self.client_id))
         self.environment.events.quitting.add_listener(on_quitting)
 

locust/test/test_main.py

@@ -133,6 +133,28 @@ def test_help_arg(self):
         self.assertIn("-f LOCUSTFILE, --locustfile LOCUSTFILE", output)
         self.assertIn("Logging options:", output)
         self.assertIn("--skip-log-setup      Disable Locust's logging setup.", output)
+
+    def test_custom_exit_code(self):
+        with temporary_file(content=textwrap.dedent("""
+            from locust import User, task, constant, events
+            @events.quitting.add_listener
+            def _(environment, **kw):
+                environment.process_exit_code = 42
+            class TestUser(User):
+                wait_time = constant(3)
+                @task
+                def my_task():
+                    print("running my_task()")
+        """)) as file_path:
+            proc = subprocess.Popen(["locust", "-f", file_path], stdout=PIPE, stderr=PIPE)
+            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)
 
     def test_webserver(self):
         with temporary_file(content=textwrap.dedent("""