locust 2.29.2.dev15__py3-none-any.whl → 2.29.2.dev32__py3-none-any.whl

locust/_version.py

@@ -12,5 +12,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '2.29.2.dev15'
-__version_tuple__ = version_tuple = (2, 29, 2, 'dev15')
+__version__ = version = '2.29.2.dev32'
+__version_tuple__ = version_tuple = (2, 29, 2, 'dev32')

locust/test/test_sequential_taskset.py

@@ -34,6 +34,26 @@ class TestTaskSet(LocustTestCase):
         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 = []
 

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 @@ class SequentialTaskSetMeta(TaskSetMeta):
                 # 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)