Merge pull request #1841 from locustio/show-custom-arguments-in-web-ui

Show custom arguments in web ui and forward them to worker

docs/configuration.rst

@@ -7,7 +7,7 @@ Configuration
 Command Line Options
 -----------------------------
 
-The most straight forward way to Configure how Locust is run is through command line arguments. 
+Locust is configured mainly through command line arguments.
 
 .. code-block:: console
 
@@ -77,6 +77,11 @@ Here's a table of all the available configuration options, and their correspondi
 
 .. include:: config-options.rst
 
+Custom arguments
+----------------
+
+See :ref:`parametrizing-locustfiles`
+
 Customization of statistics settings
 ------------------------------------
 

docs/extending-locust.rst

@@ -147,6 +147,34 @@ For example, you can monitor the fail ratio of your test and stop the run if it
         if not isinstance(environment.runner, WorkerRunner):
             gevent.spawn(checker, environment)
 
+.. _parametrizing-locustfiles:
+
+Parametrizing locustfiles
+=========================
+
+There are two main ways to parametrize your locustfile. You can use environment variables:
+
+On linux/mac:
+
+.. code-block:: bash
+
+    MY_FUNKY_VAR=42 locust ...
+
+On windows:
+
+.. code-block:: bash
+
+    SET MY_FUNKY_VAR=42
+    locust ...
+
+And then access them in your locustfile:
+
+.. code-block:: python
+
+    import os
+    print(os.environ['MY_FUNKY_VAR'])
+
+Or you can add your own custom command line arguments, using the :py:attr:`init_command_line_parser <locust.event.Events.init_command_line_parser>` event, as shown in `this example <https://github.com/locustio/locust/tree/master/examples/add_command_line_argument.py>`_. Custom arguments can also be set in the start dialogue in the web UI. When running Locust :ref:`distributed <running-locust-distributed>`, custom arguments are automatically forwarded to workers when the run is started (but not before then, so you cannot rely on forwarded arguments *before* the test has actually started).
 
 More examples
 =============

examples/add_command_line_argument.py

@@ -13,13 +13,9 @@ def _(environment, **kw):
 
 
 class WebsiteUser(HttpUser):
-    """
-    User class that does requests to the locust web server running on localhost
-    """
-
     host = "http://127.0.0.1:8089"
     wait_time = between(2, 5)
 
     @task
     def my_task(self):
-        pass
+        print(self.environment.parsed_options.my_argument)

locust/argument_parser.py

@@ -443,3 +443,10 @@ def parse_options(args=None):
     if parsed_opts.stats_history_enabled and (parsed_opts.csv_prefix is None):
         parser.error("'--csv-full-history' requires '--csv'.")
     return parsed_opts
+
+
+def default_args_dict():
+    # returns a dict containing the default arguments (before any custom arguments are added)
+    default_parser = get_empty_argument_parser()
+    setup_parser_arguments(default_parser)
+    return vars(default_parser.parse([]))

locust/runners.py

@@ -41,6 +41,7 @@
     RequestStats,
     setup_distributed_stats_event_listeners,
 )
+from . import argument_parser
 
 logger = logging.getLogger(__name__)
 
@@ -690,6 +691,9 @@ def start(self, user_count: int, spawn_rate: float, **kwargs) -> None:
                         "user_classes_count": worker_user_classes_count,
                         "host": self.environment.host,
                         "stop_timeout": self.environment.stop_timeout,
+                        "parsed_options": vars(self.environment.parsed_options)
+                        if self.environment.parsed_options
+                        else {},
                     }
                     dispatch_greenlets.add(
                         gevent.spawn_later(
@@ -1112,6 +1116,17 @@ def worker(self):
                     continue
                 self.environment.host = job["host"]
                 self.environment.stop_timeout = job["stop_timeout"]
+
+                # receive custom arguments
+                if self.environment.parsed_options is None:
+                    default_parser = argument_parser.get_empty_argument_parser()
+                    argument_parser.setup_parser_arguments(default_parser)
+                    self.environment.parsed_options = default_parser.parse(args=[])
+                custom_args_from_master = {
+                    k: v for k, v in job["parsed_options"].items() if k not in argument_parser.default_args_dict()
+                }
+                vars(self.environment.parsed_options).update(custom_args_from_master)
+
                 if self.spawning_greenlet:
                     # kill existing spawning greenlet before we launch new one
                     self.spawning_greenlet.kill(block=True)

locust/templates/index.html

@@ -78,6 +78,11 @@ <h2>Start new load test</h2>
                         {% endif %}
                     </label>
                     <input type="text" name="host" id="host" class="val" autocapitalize="off" autocorrect="off" value="{{ host  or "" }}" onfocus="this.select()"/><br>
+                    {% if extra_options %}<label>Custom parameters:</label>{% endif %}
+                    {% for key, value in extra_options.items() %}
+                        <label for="{{key}}">{{key}}</label>
+                        <input type="text" name="{{key}}" id="{{key}}" class="val" value="{{value}}" /><br>
+                    {% endfor %}
                     <button type="submit">Start swarming</button>
                 </form>
                 <div style="clear:right;"></div>

locust/test/test_runners.py

@@ -21,6 +21,7 @@
     runners,
     __version__,
 )
+from locust.argument_parser import parse_options
 from locust.env import Environment
 from locust.exception import (
     RPCError,
@@ -615,8 +616,8 @@ class TestUser(User):
             wait_time = constant(0.1)
 
             @task
-            def incr_stats(l):
-                l.environment.events.request.fire(
+            def incr_stats(self):
+                self.environment.events.request.fire(
                     request_type="GET",
                     name="/",
                     response_time=1337,
@@ -659,6 +660,69 @@ def incr_stats(l):
             "For some reason the master node's stats has not come in",
         )
 
+    def test_distributed_run_with_custom_args(self):
+        """
+        Full integration test that starts both a MasterRunner and three WorkerRunner instances
+        and makes sure that their stats is sent to the Master.
+        """
+
+        class TestUser(User):
+            wait_time = constant(0.1)
+
+            @task
+            def incr_stats(self):
+                self.environment.events.request.fire(
+                    request_type="GET",
+                    name=self.environment.parsed_options.my_str_argument,
+                    response_time=self.environment.parsed_options.my_int_argument,
+                    response_length=666,
+                    exception=None,
+                    context={},
+                )
+
+        @locust.events.init_command_line_parser.add_listener
+        def _(parser, **kw):
+            parser.add_argument("--my-int-argument", type=int)
+            parser.add_argument("--my-str-argument", type=str, default="NOOOO")
+
+        with mock.patch("locust.runners.WORKER_REPORT_INTERVAL", new=0.3):
+            # start a Master runner
+            master_env = Environment(user_classes=[TestUser])
+            master = master_env.create_master_runner("*", 0)
+            master_env.parsed_options = parse_options(
+                [
+                    "--my-int-argument",
+                    "42",
+                    "--my-str-argument",
+                    "cool-string",
+                ]
+            )
+            sleep(0)
+            # start 3 Worker runners
+            workers = []
+            for i in range(3):
+                worker_env = Environment(user_classes=[TestUser])
+                worker = worker_env.create_worker_runner("127.0.0.1", master.server.port)
+                workers.append(worker)
+
+            # give workers time to connect
+            sleep(0.1)
+            # issue start command that should trigger TestUsers to be spawned in the Workers
+            master.start(6, spawn_rate=1000)
+            sleep(0.1)
+            # check that worker nodes have started locusts
+            for worker in workers:
+                self.assertEqual(2, worker.user_count)
+            # give time for users to generate stats, and stats to be sent to master
+            sleep(1)
+            master.quit()
+            # make sure users are killed
+            for worker in workers:
+                self.assertEqual(0, worker.user_count)
+
+        self.assertEqual(master_env.runner.stats.total.max_response_time, 42)
+        self.assertEqual(master_env.runner.stats.get("cool-string", "GET").avg_response_time, 42)
+
     def test_test_stop_event(self):
         class TestUser(User):
             wait_time = constant(0.1)
@@ -2267,6 +2331,7 @@ def the_task(self):
                         "user_classes_count": {"MyTestUser": 1},
                         "host": "",
                         "stop_timeout": 1,
+                        "parsed_options": {},
                     },
                     "dummy_client_id",
                 )
@@ -2307,6 +2372,7 @@ def the_task(self):
                         "user_classes_count": {"MyTestUser": 1},
                         "host": "",
                         "stop_timeout": None,
+                        "parsed_options": {},
                     },
                     "dummy_client_id",
                 )
@@ -2351,6 +2417,7 @@ def my_task(self):
                         "user_classes_count": {"MyUser": 10},
                         "host": "",
                         "stop_timeout": None,
+                        "parsed_options": {},
                     },
                     "dummy_client_id",
                 )
@@ -2369,6 +2436,7 @@ def my_task(self):
                         "user_classes_count": {"MyUser": 9},
                         "host": "",
                         "stop_timeout": None,
+                        "parsed_options": {},
                     },
                     "dummy_client_id",
                 )
@@ -2386,6 +2454,7 @@ def my_task(self):
                         "user_classes_count": {"MyUser": 2},
                         "host": "",
                         "stop_timeout": None,
+                        "parsed_options": {},
                     },
                     "dummy_client_id",
                 )
@@ -2403,6 +2472,7 @@ def my_task(self):
                         "user_classes_count": {"MyUser": 2},
                         "host": "",
                         "stop_timeout": None,
+                        "parsed_options": {},
                     },
                     "dummy_client_id",
                 )
@@ -2443,6 +2513,7 @@ def my_task(self):
                         "user_classes_count": {"MyUser": 10},
                         "host": "",
                         "stop_timeout": None,
+                        "parsed_options": {},
                     },
                     "dummy_client_id",
                 )
@@ -2524,6 +2595,7 @@ def my_task(self):
                         "user_classes_count": {"MyUser": 10},
                         "host": "",
                         "stop_timeout": None,
+                        "parsed_options": {},
                     },
                     "dummy_client_id",
                 )
@@ -2538,6 +2610,7 @@ def my_task(self):
                         "user_classes_count": {"MyUser": 9},
                         "host": "",
                         "stop_timeout": None,
+                        "parsed_options": {},
                     },
                     "dummy_client_id",
                 )
@@ -2574,6 +2647,7 @@ def my_task(self):
                         "user_classes_count": {"MyUser1": 10, "MyUser2": 10},
                         "host": "",
                         "stop_timeout": None,
+                        "parsed_options": {},
                     },
                     "dummy_client_id",
                 )
@@ -2591,6 +2665,7 @@ def my_task(self):
                         "user_classes_count": {"MyUser1": 1, "MyUser2": 2},
                         "host": "",
                         "stop_timeout": None,
+                        "parsed_options": {},
                     },
                     "dummy_client_id",
                 )
@@ -2702,6 +2777,7 @@ def on_test_start(*args, **kw):
                         "num_users": 1,
                         "host": "",
                         "stop_timeout": None,
+                        "parsed_options": {},
                     },
                     "dummy_client_id",
                 )
@@ -2727,6 +2803,7 @@ def on_test_start(*args, **kw):
                         "num_users": 1,
                         "host": "",
                         "stop_timeout": None,
+                        "parsed_options": {},
                     },
                     "dummy_client_id",
                 )
@@ -2745,6 +2822,7 @@ def on_test_start(*args, **kw):
                         "num_users": 1,
                         "host": "",
                         "stop_timeout": None,
+                        "parsed_options": {},
                     },
                     "dummy_client_id",
                 )
@@ -2785,6 +2863,7 @@ def on_test_stop(*args, **kw):
                         "num_users": 1,
                         "host": "",
                         "stop_timeout": None,
+                        "parsed_options": {},
                     },
                     "dummy_client_id",
                 )
@@ -2820,6 +2899,7 @@ def on_test_stop(*args, **kw):
                         "num_users": 1,
                         "host": "",
                         "stop_timeout": None,
+                        "parsed_options": {},
                     },
                     "dummy_client_id",
                 )

locust/test/test_web.py

@@ -242,6 +242,30 @@ def my_task(self):
         self.assertEqual("https://localhost", response.json()["host"])
         self.assertEqual(self.environment.host, "https://localhost")
 
+    def test_swarm_custom_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, **kw):
+            parser.add_argument("--my-argument", type=int, help="Give me a number")
+
+        self.environment.user_classes = [MyUser]
+        self.environment.parsed_options = parse_options(args=["--my-argument", "42"])
+        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(my_dict["val"], 42)
+
     def test_swarm_host_value_not_specified(self):
         class MyUser(User):
             wait_time = constant(1)