diff --git a/locust/test/test_web.py b/locust/test/test_web.py index 67b1be3fed..93613e8112 100644 --- a/locust/test/test_web.py +++ b/locust/test/test_web.py @@ -764,7 +764,7 @@ def tick(self): response = requests.get("http://127.0.0.1:%i/stop" % self.web_port) self.assertEqual(response.json()["message"], "Test stopped") - def test_swarm_custom_argument(self): + def test_swarm_custom_argument_without_default_value(self): my_dict = {} class MyUser(User): @@ -776,10 +776,10 @@ def my_task(self): my_dict["val"] = self.environment.parsed_options.my_argument @locust.events.init_command_line_parser.add_listener - def _(parser, **kw): + def _(parser): parser.add_argument("--my-argument", type=int, help="Give me a number") - parsed_options = parse_options(args=["--my-argument", "42"]) + parsed_options = parse_options() self.environment.user_classes = [MyUser] self.environment.parsed_options = parsed_options self.environment.web_ui.parsed_options = parsed_options @@ -788,7 +788,59 @@ def _(parser, **kw): data={"user_count": 1, "spawn_rate": 1, "host": "", "my_argument": "42"}, ) self.assertEqual(200, response.status_code) - self.assertEqual(my_dict["val"], 42) + self.assertEqual("42", my_dict["val"]) + + def test_swarm_custom_argument_with_default_value(self): + my_dict = {} + + class MyUser(User): + host = "http://example.com" + wait_time = constant(1) + + @task(1) + def my_task(self): + my_dict["val"] = self.environment.parsed_options.my_argument + + @locust.events.init_command_line_parser.add_listener + def _(parser): + parser.add_argument("--my-argument", type=int, help="Give me a number", default=24) + + parsed_options = parse_options() + self.environment.user_classes = [MyUser] + self.environment.parsed_options = parsed_options + self.environment.web_ui.parsed_options = parsed_options + response = requests.post( + "http://127.0.0.1:%i/swarm" % self.web_port, + data={"user_count": 1, "spawn_rate": 1, "host": "", "my_argument": "42"}, + ) + self.assertEqual(200, response.status_code) + self.assertEqual(42, my_dict["val"]) + + def test_swarm_override_command_line_argument(self): + my_dict = {} + + class MyUser(User): + host = "http://example.com" + wait_time = constant(1) + + @task(1) + def my_task(self): + my_dict["val"] = self.environment.parsed_options.my_argument + + @locust.events.init_command_line_parser.add_listener + def _(parser): + parser.add_argument("--my-argument", type=int, help="Give me a number") + + parsed_options = parse_options(args=["--my-argument", "24"]) + self.environment.user_classes = [MyUser] + self.environment.parsed_options = parsed_options + self.environment.web_ui.parsed_options = parsed_options + response = requests.post( + "http://127.0.0.1:%i/swarm" % self.web_port, + data={"user_count": 1, "spawn_rate": 1, "host": "", "my_argument": "42"}, + ) + self.assertEqual(200, response.status_code) + self.assertEqual(42, my_dict["val"]) def test_swarm_host_value_not_specified(self): class MyUser(User): diff --git a/locust/web.py b/locust/web.py index 90c41fcde2..ccbfe571e4 100644 --- a/locust/web.py +++ b/locust/web.py @@ -285,9 +285,9 @@ def swarm() -> Response: if isinstance(parsed_options_value, bool): parsed_options_dict[key] = value == "true" elif parsed_options_value is None: - parsed_options_dict[key] = parsed_options_value + parsed_options_dict[key] = value else: - parsed_options_dict[key] = type(parsed_options_dict[key])(value) + parsed_options_dict[key] = type(parsed_options_value)(value) if environment.shape_class and environment.runner is not None: environment.runner.start_shape()