diff --git a/locust/argument_parser.py b/locust/argument_parser.py index d7f617a7e9..22817a898e 100644 --- a/locust/argument_parser.py +++ b/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( diff --git a/locust/env.py b/locust/env.py index 05b9e64e13..51114de0f5 100644 --- a/locust/env.py +++ b/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} diff --git a/locust/main.py b/locust/main.py index 01fea3dd94..1288d1e556 100644 --- a/locust/main.py +++ b/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) diff --git a/locust/test/test_env.py b/locust/test/test_env.py index 099c7ee6b1..dccaad1527 100644 --- a/locust/test/test_env.py +++ b/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"])