Add -i parameter to provide a workaround for locustio/locust#1085

Also move the checks parameters into __init__.py

README.md

@@ -14,7 +14,11 @@ The plugins are grouped by type:
 * readers (ways to get test data into your tests) - currently implemented [CSV](https://github.com/SvenskaSpel/locust-plugins/blob/master/locust_plugins/csvreader.py) and [MongoDB](https://github.com/SvenskaSpel/locust-plugins/blob/master/locust_plugins/mongoreader.py)
 * [wait time](locust_plugins/wait_time.py) (custom wait time functions)
 * [debug](locust_plugins/debug.py) (support for running a single user in the debugger)
-* [checks](locust_plugins/checks.py) (adds command line parameters to set locust exit code based on requests/s, error percentage and average response times)
+
+* [Command line options](locust_plugins/__init__.py)
+    - Iteration limit (`-i`), stops Locust after a certain number of task iterations
+    - Checks (`--check-rps`, `--check-fail-ratio`, `--check-avg-response-time`), gives an error return code if certain conditions are not met
+    - Here are some [examples of use](examples/cmd_line_examples.sh)
 
 Have a look at the [example locustfiles](examples/) to learn how to use the plugins.
 

examples/cmd_line_examples.sh

@@ -0,0 +1,13 @@
+# for command line options to work, you need to run a locust file that imports locust_plugins
+
+locust -u 1 -t 60 --headless --check-rps 5 --check-fail-ratio 0.05 --check-avg-response-time 50
+# Set locust's exit code to failed (2) if any of the following are not met:
+# * At least 5 requests/s
+# * At most 5% errors
+# * At most 50ms average response times
+# (all values are for the whole run)
+
+locust -u 5 -t 60 --headless -i 10
+# Stop locust after 10 task iterations (this is an upper bound, so you can be sure no more than 10 of iterations will be done)
+# Note that the limit is applied *per worker* in a distributed run. So if you, for example, have 2 workers you will run 20 requests.
+# It is (currently) not distributed from master to worker.

locust_plugins/__init__.py

@@ -1,10 +1,97 @@
 __version__ = "1.0.14"
-import time
 from .wait_time import constant_ips, constant_total_ips
 from .debug import run_single_user
-from locust import User, constant
-from . import checks as _checks  # import this for the side-effects
+from locust import User, constant, TaskSet
+from locust import events
+from locust.exception import StopUser
+from locust.env import Environment
+from locust.runners import Runner
+import logging
 
 # Monkey patch User while waiting for everyone else to see the light:
 # https://github.com/locustio/locust/issues/1308
 User.wait_time = constant(0)
+
+
+@events.init_command_line_parser.add_listener
+def add_checks_arguments(parser):
+    checks = parser.add_argument_group(
+        "Checks", "Sets locust's exit code to 2 if any of these thresholds were not met (added by locust-plugins)"
+    )
+    checks.add_argument("--check-rps", type=float, help="Requests per second", env_var="LOCUST_CHECK_RPS", default=0.0)
+    checks.add_argument(
+        "--check-fail-ratio",
+        type=float,
+        help="Ratio of failed requests (0.0-1.0)",
+        env_var="LOCUST_CHECK_FAIL_RATIO",
+        default=1.0,
+    )
+    checks.add_argument(
+        "--check-avg-response-time",
+        type=float,
+        help="Average response time",
+        env_var="LOCUST_CHECK_AVG_RESPONSE_TIME",
+        default=float("inf"),
+    )
+    other = parser.add_argument_group("Plugins", "Other extra parameters added by locust-plugins")
+    # fix for https://github.com/locustio/locust/issues/1085
+    other.add_argument(
+        "-i",
+        "--iterations",
+        type=int,
+        help="Dont run more than this number of task iterations and terminate once they have finished",
+        env_var="LOCUST_ITERATIONS",
+        default=0,
+    )
+
+
+@events.test_start.add_listener
+def set_up_iteration_limit(environment: Environment, **_kwargs):
+    if environment.parsed_options.iterations:
+        runner: Runner = environment.runner
+        runner.iterations_started = 0
+        runner.iteration_target_reached = False
+
+        # monkey patch Runner to add support for iterations limit
+        _execute_next_task = TaskSet.execute_next_task
+
+        def execute_next_task_with_iteration_limit(self: TaskSet):
+            if runner.iterations_started == environment.parsed_options.iterations:
+                if not runner.iteration_target_reached:
+                    runner.iteration_target_reached = True
+                    logging.info(
+                        f"Iteration limit reached ({environment.parsed_options.iterations}), stopping Users at the start of their next task run"
+                    )
+                if runner.user_count == 1:
+                    logging.debug("Last user stopped, quitting runner")
+                    runner.quit()
+                raise StopUser()
+            runner.iterations_started = runner.iterations_started + 1
+            _execute_next_task(self)
+
+        TaskSet.execute_next_task = execute_next_task_with_iteration_limit
+
+
+@events.quitting.add_listener
+def do_checks(environment, **_kw):
+    stats = environment.runner.stats.total
+    fail_ratio = stats.fail_ratio
+    total_rps = stats.total_rps
+    avg_response_time = stats.avg_response_time
+
+    opts = environment.parsed_options
+    check_rps = opts.check_rps
+    check_fail_ratio = opts.check_fail_ratio
+    check_avg_response_time = opts.check_avg_response_time
+
+    if fail_ratio > check_fail_ratio:
+        logging.info(f"Check failed: fail ratio was {fail_ratio:.1f} (threshold {check_fail_ratio:.1f})")
+        environment.process_exit_code = 2
+    if total_rps < check_rps:
+        logging.info(f"Check failed: total rps was {total_rps:.1f} (threshold {check_rps:.1f})")
+        environment.process_exit_code = 2
+    if avg_response_time > check_avg_response_time:
+        logging.info(
+            f"Check failed: avg response time was {avg_response_time:.1f} (threshold {check_avg_response_time:.1f})"
+        )
+        environment.process_exit_code = 2