Add equal weights flag

locust/argument_parser.py

@@ -417,6 +417,13 @@ def setup_parser_arguments(parser):
         help="Number of seconds to wait for a simulated user to complete any executing task before exiting. Default is to terminate immediately. This parameter only needs to be specified for the master process when running Locust distributed.",
         env_var="LOCUST_STOP_TIMEOUT",
     )
+    other_group.add_argument(
+        "--equal-weights",
+        action="store_true",
+        default=False,
+        dest="equal_weights",
+        help="Use equally distributed task weights, overriding the weights specified in the locustfile.",
+    )
 
     user_classes_group = parser.add_argument_group("User classes")
     user_classes_group.add_argument(

locust/env.py

@@ -1,5 +1,6 @@
 from operator import methodcaller
 from typing import (
+    Callable,
     Dict,
     List,
     Type,
@@ -12,7 +13,7 @@
 from .runners import Runner, LocalRunner, MasterRunner, WorkerRunner
 from .web import WebUI
 from .user import User
-from .user.task import filter_tasks_by_tags
+from .user.task import TaskSet, filter_tasks_by_tags
 from .shape import LoadTestShape
 
 
@@ -215,6 +216,26 @@ def _filter_tasks_by_tags(self):
         for user_class in self.user_classes:
             filter_tasks_by_tags(user_class, self.tags, self.exclude_tags)
 
+    def assign_equal_weights(self):
+        """
+        Update the user classes such that each user runs their specified tasks with equal
+        probability.
+        """
+        for u in self.user_classes:
+            u.weight = 1
+            user_tasks = []
+            tasks_frontier = u.tasks
+            while len(tasks_frontier) != 0:
+                t = tasks_frontier.pop()
+                if hasattr(t, "tasks") and t.tasks:
+                    tasks_frontier.extend(t.tasks)
+                elif isinstance(t, Callable):
+                    if t not in user_tasks:
+                        user_tasks.append(t)
+                else:
+                    raise ValueError("Unrecognized task type in user")
+            u.tasks = user_tasks
+
     @property
     def user_classes_by_name(self) -> Dict[str, Type[User]]:
         return {u.__name__: u for u in self.user_classes}

locust/main.py

@@ -298,6 +298,12 @@ def main():
     else:
         web_ui = None
 
+    def assign_equal_weights(environment, **kwargs):
+        environment.assign_equal_weights()
+
+    if options.equal_weights:
+        environment.events.init.add_listener(assign_equal_weights)
+
     # Fire locust init event which can be used by end-users' code to run setup code that
     # need access to the Environment, Runner or WebUI.
     environment.events.init.fire(environment=environment, runner=runner, web_ui=web_ui)

locust/test/test_env.py

@@ -6,6 +6,7 @@
     User,
     task,
 )
+from locust.user.task import TaskSet
 from .testcases import LocustTestCase
 from .fake_module1_for_env_test import MyUserWithSameName as MyUserWithSameName1
 from .fake_module2_for_env_test import MyUserWithSameName as MyUserWithSameName2
@@ -39,3 +40,106 @@ def test_user_classes_with_same_name_is_error(self):
             e.exception.args[0],
             "The following user classes have the same class name: locust.test.fake_module1_for_env_test.MyUserWithSameName, locust.test.fake_module2_for_env_test.MyUserWithSameName",
         )
+
+    def test_assign_equal_weights(self):
+        def verify_tasks(u, target_tasks):
+            self.assertEqual(len(u.tasks), len(target_tasks))
+            tasks = [t.__name__ for t in u.tasks]
+            self.assertEqual(len(tasks), len(set(tasks)))
+            self.assertEqual(set(tasks), set(target_tasks))
+
+        # Base case
+        class MyUser1(User):
+            wait_time = constant(0)
+
+            @task(4)
+            def my_task(self):
+                pass
+
+            @task(1)
+            def my_task_2(self):
+                pass
+
+        environment = Environment(user_classes=[MyUser1])
+        environment.assign_equal_weights()
+        u = environment.user_classes[0]
+        verify_tasks(u, ["my_task", "my_task_2"])
+
+        # Testing nested task sets
+        class MyUser2(User):
+            @task
+            class TopLevelTaskSet(TaskSet):
+                @task
+                class IndexTaskSet(TaskSet):
+                    @task(10)
+                    def index(self):
+                        self.client.get("/")
+
+                    @task
+                    def stop(self):
+                        self.client.get("/hi")
+
+                @task(2)
+                def stats(self):
+                    self.client.get("/stats/requests")
+
+        environment = Environment(user_classes=[MyUser2])
+        environment.assign_equal_weights()
+        u = environment.user_classes[0]
+        verify_tasks(u, ["index", "stop", "stats"])
+
+        # Testing task assignment via instance variable
+        def outside_task():
+            pass
+
+        def outside_task_2():
+            pass
+
+        class SingleTaskSet(TaskSet):
+            tasks = [outside_task, outside_task, outside_task_2]
+
+        class MyUser3(User):
+            tasks = [SingleTaskSet, outside_task]
+
+        environment = Environment(user_classes=[MyUser3])
+        environment.assign_equal_weights()
+        u = environment.user_classes[0]
+        verify_tasks(u, ["outside_task", "outside_task_2"])
+
+        # Testing task assignment via dict
+        class DictTaskSet(TaskSet):
+            def dict_task_1():
+                pass
+
+            def dict_task_2():
+                pass
+
+            def dict_task_3():
+                pass
+
+            tasks = {
+                dict_task_1: 5,
+                dict_task_2: 3,
+                dict_task_3: 1,
+            }
+
+        class MyUser4(User):
+            tasks = [DictTaskSet, SingleTaskSet, SingleTaskSet]
+
+        # Assign user tasks in dict
+        environment = Environment(user_classes=[MyUser4])
+        environment.assign_equal_weights()
+        u = environment.user_classes[0]
+        verify_tasks(u, ["outside_task", "outside_task_2", "dict_task_1", "dict_task_2", "dict_task_3"])
+
+        class MyUser5(User):
+            tasks = {
+                DictTaskSet: 5,
+                SingleTaskSet: 3,
+                outside_task: 6,
+            }
+
+        environment = Environment(user_classes=[MyUser5])
+        environment.assign_equal_weights()
+        u = environment.user_classes[0]
+        verify_tasks(u, ["outside_task", "outside_task_2", "dict_task_1", "dict_task_2", "dict_task_3"])