diff --git a/locust/argument_parser.py b/locust/argument_parser.py index 3a283953f9..fb6144acbc 100644 --- a/locust/argument_parser.py +++ b/locust/argument_parser.py @@ -583,7 +583,8 @@ def setup_parser_arguments(parser): "user_classes", nargs="*", metavar="UserClass", - help="Optionally specify which User classes that should be used (available User classes can be listed with -l or --list)", + help="Optionally specify which User classes that should be used (available User classes can be listed with -l or --list). LOCUST_USER_CLASSES environment variable can also be used to specify User classes", + default=os.environ.get("LOCUST_USER_CLASSES", "").split(), ) diff --git a/locust/test/test_parser.py b/locust/test/test_parser.py index dea9bec533..c198ac5b1e 100644 --- a/locust/test/test_parser.py +++ b/locust/test/test_parser.py @@ -93,6 +93,34 @@ def test_parse_options(self): # check default arg self.assertEqual(8089, options.web_port) + def test_parse_options_from_env(self): + os.environ["LOCUST_LOCUSTFILE"] = "locustfile.py" + os.environ["LOCUST_USERS"] = "100" + os.environ["LOCUST_SPAWN_RATE"] = "10" + os.environ["LOCUST_RUN_TIME"] = "5m" + os.environ["LOCUST_RESET_STATS"] = "true" + os.environ["LOCUST_STOP_TIMEOUT"] = "5" + os.environ["LOCUST_USER_CLASSES"] = "MyUserClass" + options = parse_options(args=[]) + + self.assertEqual("locustfile.py", options.locustfile) + self.assertEqual(100, options.num_users) + self.assertEqual(10, options.spawn_rate) + self.assertEqual("5m", options.run_time) + self.assertTrue(options.reset_stats) + self.assertEqual("5", options.stop_timeout) + self.assertEqual(["MyUserClass"], options.user_classes) + # check default arg + self.assertEqual(8089, options.web_port) + + del os.environ["LOCUST_LOCUSTFILE"] + del os.environ["LOCUST_USERS"] + del os.environ["LOCUST_SPAWN_RATE"] + del os.environ["LOCUST_RUN_TIME"] + del os.environ["LOCUST_RESET_STATS"] + del os.environ["LOCUST_STOP_TIMEOUT"] + del os.environ["LOCUST_USER_CLASSES"] + def test_parse_locustfile(self): with mock_locustfile() as mocked: locustfiles = parse_locustfile_option(