Merge pull request #1310 from locustio/sequential-taskset

Replace TaskSequence and @seq_task with SequentialTaskSet

docs/api.rst

@@ -27,16 +27,12 @@ task decorator
 
 .. autofunction:: locust.core.task
 
-TaskSequence class
-==================
+SequentialTaskSet class
+=======================
 
-.. autoclass:: locust.core.TaskSequence
+.. autoclass:: locust.sequential_taskset.SequentialTaskSet
     :members: locust, parent, wait_time, client, tasks, interrupt, schedule_task, on_start, on_stop
 
-seq_task decorator
-==================
-
-.. autofunction:: locust.core.seq_task
 
 .. _wait_time_functions:
 

docs/writing-a-locustfile.rst

@@ -320,30 +320,42 @@ parent TaskSet instance.
 
 
 
-TaskSequence class
-==================
+SequentialTaskSet class
+=======================
 
-TaskSequence class is a TaskSet but its tasks will be executed in order.
-To define this order you should do the following:
+:py:class:`SequentialTaskSet <locust.sequential_taskset.SequentialTaskSet>` is a TaskSet but its 
+tasks will be executed in the order that they are declared. Weights are ignored for tasks on a 
+SequentialTaskSet class. Ofcourse you can also nest SequentialTaskSet within TaskSet and vice versa.
 
 .. code-block:: python
-
-    class MyTaskSequence(TaskSequence):
-        @seq_task(1)
+    
+    def function_task(taskset):
+        pass
+    
+    class SequenceOfTasks(SequentialTaskSet):
+        @task
         def first_task(self):
             pass
-
-        @seq_task(2)
+        
+        tasks = [functon_task]
+        
+        @task
         def second_task(self):
             pass
 
-        @seq_task(3)
-        @task(10)
+        @task
         def third_task(self):
             pass
 
-In the above example, the order is defined to execute first_task, then second_task and lastly the third_task for 10 times.
-As you can see, you can compose :py:meth:`@seq_task <locust.core.seq_task>` with :py:meth:`@task <locust.core.task>` decorator, and of course you can also nest TaskSets within TaskSequences and vice versa.
+In the above example, the tasks are executed in the order of declaration: 
+
+1. ``first_task``
+2. ``function_task``
+3. ``second_task``
+4. ``third_task``
+
+and then it will start over at ``first_task`` again.
+
 
 .. _on-start-on-stop:
 

examples/browse_docs_sequence_test.py

@@ -2,16 +2,16 @@
 # browsing the Locust documentation on https://docs.locust.io/
 
 import random
-from locust import HttpLocust, TaskSequence, seq_task, task, between
+from locust import HttpLocust, SequentialTaskSet, task, between
 from pyquery import PyQuery
 
 
-class BrowseDocumentationSequence(TaskSequence):
+class BrowseDocumentationSequence(SequentialTaskSet):
     def on_start(self):
         self.urls_on_current_page = self.toc_urls = None
 
     # assume all users arrive at the index page
-    @seq_task(1)
+    @task
     def index_page(self):
         r = self.client.get("/")
         pq = PyQuery(r.content)
@@ -20,8 +20,7 @@ def index_page(self):
             l.attrib["href"] for l in link_elements
         ]
 
-    @seq_task(2)
-    @task(50)
+    @task
     def load_page(self, url=None):
         url = random.choice(self.toc_urls)
         r = self.client.get(url)
@@ -31,8 +30,7 @@ def load_page(self, url=None):
             l.attrib["href"] for l in link_elements
         ]
 
-    @seq_task(3)
-    @task(30)
+    @task
     def load_sub_page(self):
         url = random.choice(self.urls_on_current_page)
         r = self.client.get(url)

locust/__init__.py

@@ -1,7 +1,9 @@
-from .core import HttpLocust, Locust, TaskSet, TaskSequence, task, seq_task
+from .core import HttpLocust, Locust, TaskSet, task
+from .event import Events
 from .exception import InterruptTaskSet, ResponseError, RescheduleTaskImmediately
+from .sequential_taskset import SequentialTaskSet
 from .wait_time import between, constant, constant_pacing
-from .event import Events
+
 events = Events()
 
 __version__ = "0.14.5"

locust/core.py

@@ -58,35 +58,6 @@ def my_task()
         return decorator_func
 
 
-def seq_task(order):
-    """
-    Used as a convenience decorator to be able to declare tasks for a TaskSequence
-    inline in the class. Example::
-
-        class NormalUser(TaskSequence):
-            @seq_task(1)
-            def login_first(self):
-                pass
-
-            @seq_task(2)
-            @task(25) # You can also set the weight in order to execute the task for `weight` times one after another.
-            def then_read_thread(self):
-                pass
-
-            @seq_task(3)
-            def then_logout(self):
-                pass
-    """
-
-    def decorator_func(func):
-        func.locust_task_order = order
-        if not hasattr(func, 'locust_task_weight'):
-            func.locust_task_weight = 1
-        return func
-
-    return decorator_func
-
-
 class NoClientWarningRaiser(object):
     """
     The purpose of this class is to emit a sensible error message for old test scripts that 
@@ -379,37 +350,6 @@ def client(self):
         return self.locust.client
 
 
-class TaskSequence(TaskSet):
-    """
-    Class defining a sequence of tasks that a Locust user will execute.
-
-    When a TaskSequence starts running, it will pick the task in `index` from the *tasks* attribute,
-    execute it, and call its *wait_function* which will define a time to sleep for.
-    This defaults to a uniformly distributed random number between *min_wait* and
-    *max_wait* milliseconds. It will then schedule the `index + 1 % len(tasks)` task for execution and so on.
-
-    TaskSequence can be nested with TaskSet, which means that a TaskSequence's *tasks* attribute can contain
-    TaskSet instances as well as other TaskSequence instances. If the nested TaskSet is scheduled to be executed, it will be
-    instantiated and called from the current executing TaskSet. Execution in the
-    currently running TaskSet will then be handed over to the nested TaskSet which will
-    continue to run until it throws an InterruptTaskSet exception, which is done when
-    :py:meth:`TaskSet.interrupt() <locust.core.TaskSet.interrupt>` is called. (execution
-    will then continue in the first TaskSet).
-
-    In this class, tasks should be defined as a list, or simply define the tasks with the @seq_task decorator
-    """
-
-    def __init__(self, parent):
-        super(TaskSequence, self).__init__(parent)
-        self._index = 0
-        self.tasks.sort(key=lambda t: t.locust_task_order if hasattr(t, 'locust_task_order') else 1)
-
-    def get_next_task(self):
-        task = self.tasks[self._index]
-        self._index = (self._index + 1) % len(self.tasks)
-        return task
-
-
 class DefaultTaskSet(TaskSet):
     """
     Default root TaskSet that executes tasks in Locust.tasks.

locust/sequential_taskset.py

@@ -0,0 +1,56 @@
+from .core import TaskSet, TaskSetMeta
+from .exception import LocustError
+
+
+class SequentialTaskSetMeta(TaskSetMeta):
+    """
+    Meta class for SequentialTaskSet. It's used to allow SequentialTaskSet classes to specify 
+    task execution in both a list as the tasks attribute or using the @task decorator
+    
+    We use the fact that class_dict order is the order of declaration in Python 3.6
+    (See https://www.python.org/dev/peps/pep-0520/)
+    """
+    def __new__(mcs, classname, bases, class_dict):
+        new_tasks = []
+        for base in bases:
+            # first get tasks from base classes
+            if hasattr(base, "tasks") and base.tasks:
+                new_tasks += base.tasks
+        for key, value in class_dict.items():
+            if key == "tasks":
+                # we want to insert tasks from the tasks attribute at the point of it's declaration 
+                # compared to methods declared with @task
+                if isinstance(value, list):
+                    new_tasks.extend(value)
+                else:
+                    raise ValueError("On SequentialTaskSet the task attribute can only be set to a list")
+
+            if "locust_task_weight" in dir(value):
+                # method decorated with @task
+                new_tasks.append(value)
+
+        class_dict["tasks"] = new_tasks
+        return type.__new__(mcs, classname, bases, class_dict)
+
+
+class SequentialTaskSet(TaskSet, metaclass=SequentialTaskSetMeta):
+    """
+    Class defining a sequence of tasks that a Locust user will execute.
+    
+    Works like TaskSet, but task weight is ignored, and all tasks are executed in order. Tasks can 
+    either be specified by setting the *tasks* attribute to a list of tasks, or by declaring tasks 
+    as methods using the @task decorator. The order of declaration decides the order of execution. 
+    
+    It's possible to combine a task list in the *tasks* attribute, with some tasks declared using 
+    the @task decorator. The order of declaration is respected also in that case.
+    """
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self._task_index = 0
+
+    def get_next_task(self):
+        if not self.tasks:
+            raise LocustError("No tasks defined. use the @task decorator or set the tasks property of the SequentialTaskSet")
+        task = self.tasks[self._task_index % len(self.tasks)]
+        self._task_index += 1
+        return task

locust/test/test_sequential_taskset.py

@@ -0,0 +1,120 @@
+from locust.core import Locust, task
+from locust.sequential_taskset import SequentialTaskSet
+from locust.exception import RescheduleTask
+from locust.wait_time import constant
+from .testcases import LocustTestCase
+
+
+class TestTaskSet(LocustTestCase):
+    def setUp(self):
+        super(TestTaskSet, self).setUp()
+        class User(Locust):
+            host = "127.0.0.1"
+            wait_time = constant(0)
+            _catch_exceptions = False
+        self.locust = User(self.environment)
+
+    def test_task_sequence_with_list(self):
+        log = []
+
+        def t1(ts):
+            log.append(1)
+        def t2(ts):
+            log.append(2)
+        def t3(ts):
+            log.append(3)
+            ts.interrupt(reschedule=False)
+
+        class MyTaskSequence(SequentialTaskSet):
+            tasks = [t1, t2, t3]
+
+        l = MyTaskSequence(self.locust)
+        self.assertRaises(RescheduleTask, lambda: l.run())
+        self.assertEqual([1,2,3], log)
+
+    def test_task_sequence_with_methods(self):
+        log = []
+        class MyTaskSequence(SequentialTaskSet):
+            @task
+            def t1(self):
+                log.append(1)
+            @task
+            def t2(self):
+                log.append(2)
+            @task(1)
+            def t3(self):
+                log.append(3)
+                self.interrupt(reschedule=False)
+
+        l = MyTaskSequence(self.locust)
+        self.assertRaises(RescheduleTask, lambda: l.run())
+        self.assertEqual([1,2,3], log)
+
+    def test_task_sequence_with_methods_and_list(self):
+        log = []
+        def func_t1(ts):
+            log.append(101)
+        def func_t2(ts):
+            log.append(102)
+        class MyTaskSequence(SequentialTaskSet):
+            @task
+            def t1(self):
+                log.append(1)
+            @task
+            def t2(self):
+                log.append(2)
+
+            tasks = [func_t1, func_t2]
+
+            @task(1)
+            def t3(self):
+                log.append(3)
+                self.interrupt(reschedule=False)
+
+        l = MyTaskSequence(self.locust)
+        self.assertRaises(RescheduleTask, lambda: l.run())
+        self.assertEqual([1,2, 101,102,3], log)
+
+    def test_task_sequence_with_inheritance(self):
+        log = []
+        class TS1(SequentialTaskSet):
+            @task
+            def t1(self):
+                log.append(1)
+            tasks = [lambda ts: log.append(30)]
+        class TS2(TS1):
+            tasks = [lambda ts: log.append(20)]
+            @task
+            def t2(self):
+                log.append(2)
+        class TS3(TS2):
+            @task
+            def t3(self):
+                log.append(3)
+                self.interrupt(reschedule=False)
+
+        l = TS3(self.locust)
+        self.assertRaises(RescheduleTask, lambda: l.run())
+        self.assertEqual([1,30,20,2,3], log)
+
+    def test_task_sequence_multiple_iterations(self):
+        log = []
+        class TS(SequentialTaskSet):
+            iteration_count = 0
+            @task
+            def t1(self):
+                log.append(1)
+            @task
+            def t2(self):
+                log.append(2)
+            @task(1)
+            def t3(self):
+                log.append(3)
+                self.iteration_count += 1
+                if self.iteration_count == 3:
+                    self.interrupt(reschedule=False)
+
+        l = TS(self.locust)
+        self.assertRaises(RescheduleTask, lambda: l.run())
+        self.assertEqual([1,2,3,1,2,3,1,2,3], log)
+