diff --git a/docs/tasksets.rst b/docs/tasksets.rst index 6398ef2cd9..6cecd2cfc4 100644 --- a/docs/tasksets.rst +++ b/docs/tasksets.rst @@ -148,7 +148,8 @@ SequentialTaskSet class ======================= :py:class:`SequentialTaskSet ` is a TaskSet whose tasks will be executed -in the order that they are declared. It is possible to nest SequentialTaskSets +in the order that they are declared. If the ** dictionary is used, each task will be repeated by the amount +given by the integer at the point of declaration. It is possible to nest SequentialTaskSets within a TaskSet and vice versa. For example, the following code will request URLs /1-/4 in order, and then repeat. @@ -164,8 +165,9 @@ For example, the following code will request URLs /1-/4 in order, and then repea self.client.get("/1") self.client.get("/2") - # you can still use the tasks attribute to specify a list of tasks + # you can still use the tasks attribute to specify a list or dict of tasks tasks = [function_task] + # tasks = {function_task: 1} @task def last_task(self): diff --git a/locust/test/test_sequential_taskset.py b/locust/test/test_sequential_taskset.py index 2aea5211a2..c8c6440f66 100644 --- a/locust/test/test_sequential_taskset.py +++ b/locust/test/test_sequential_taskset.py @@ -34,6 +34,26 @@ class MyTaskSequence(SequentialTaskSet): self.assertRaises(RescheduleTask, lambda: l.run()) self.assertEqual([1, 2, 3], log) + def test_task_sequence_with_dictionary(self): + log = [] + + def t1(self): + log.append(1) + + def t2(self): + log.append(2) + + def t3(self): + log.append(3) + self.interrupt(reschedule=False) + + class MyTaskSequence(SequentialTaskSet): + tasks = {t1: 3, t2: 2, t3: 1} + + l = MyTaskSequence(self.locust) + self.assertRaises(RescheduleTask, lambda: l.run()) + self.assertEqual([1, 1, 1, 2, 2, 3], log) + def test_task_sequence_with_methods(self): log = [] diff --git a/locust/user/sequential_taskset.py b/locust/user/sequential_taskset.py index 6aa0f3cf10..8a64752ea1 100644 --- a/locust/user/sequential_taskset.py +++ b/locust/user/sequential_taskset.py @@ -1,6 +1,6 @@ from locust.exception import LocustError -import logging +from itertools import cycle from .task import TaskSet, TaskSetMeta @@ -26,8 +26,12 @@ def __new__(mcs, classname, bases, class_dict): # compared to methods declared with @task if isinstance(value, list): new_tasks.extend(value) + elif isinstance(value, dict): + for task, weight in value.items(): + for _ in range(weight): + new_tasks.append(task) else: - raise ValueError("On SequentialTaskSet the task attribute can only be set to a list") + raise ValueError("The 'tasks' attribute can only be set to list or dict") if "locust_task_weight" in dir(value): # method decorated with @task @@ -52,13 +56,11 @@ class SequentialTaskSet(TaskSet, metaclass=SequentialTaskSetMeta): def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) - self._task_index = 0 + self._task_cycle = cycle(self.tasks) def get_next_task(self): if not self.tasks: raise LocustError( "No tasks defined. Use the @task decorator or set the 'tasks' attribute of the SequentialTaskSet" ) - task = self.tasks[self._task_index % len(self.tasks)] - self._task_index += 1 - return task + return next(self._task_cycle)