diff --git a/locust/env.py b/locust/env.py index eeadcbe4bc..15770d4a4b 100644 --- a/locust/env.py +++ b/locust/env.py @@ -106,6 +106,8 @@ def __init__( self._filter_tasks_by_tags() + self._remove_user_classes_with_weight_zero() + # Validate there's no class with the same name but in different modules if len(set(user_class.__name__ for user_class in self.user_classes)) != len(self.user_classes): raise ValueError( @@ -216,6 +218,19 @@ 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 _remove_user_classes_with_weight_zero(self): + """ + Remove user classes having a weight of zero. + """ + if len(self.user_classes) == 0: + # Preserve previous behaviour that allowed no user classes to be specified. + return + filtered_user_classes = [user_class for user_class in self.user_classes if user_class.weight > 0] + if len(filtered_user_classes) == 0: + # TODO: Better exception than `ValueError`? + raise ValueError("There are no users with weight > 0.") + self.user_classes[:] = filtered_user_classes + def assign_equal_weights(self): """ Update the user classes such that each user runs their specified tasks with equal diff --git a/locust/test/test_env.py b/locust/test/test_env.py index dccaad1527..09ba4e97a3 100644 --- a/locust/test/test_env.py +++ b/locust/test/test_env.py @@ -143,3 +143,50 @@ class MyUser5(User): 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"]) + + def test_user_classes_with_zero_weight_are_removed(self): + class MyUser1(User): + wait_time = constant(0) + weight = 0 + + @task + def my_task(self): + pass + + class MyUser2(User): + wait_time = constant(0) + weight = 1 + + @task + def my_task(self): + pass + + environment = Environment(user_classes=[MyUser1, MyUser2]) + + self.assertEqual(len(environment.user_classes), 1) + self.assertIs(environment.user_classes[0], MyUser2) + + def test_all_user_classes_with_zero_weight_raises_exception(self): + class MyUser1(User): + wait_time = constant(0) + weight = 0 + + @task + def my_task(self): + pass + + class MyUser2(User): + wait_time = constant(0) + weight = 0 + + @task + def my_task(self): + pass + + with self.assertRaises(ValueError) as e: + environment = Environment(user_classes=[MyUser1, MyUser2]) + + self.assertEqual( + e.exception.args[0], + "There are no users with weight > 0.", + )