Add tasks sequence support (#827)

* Add support for tasks sequence

* Add TaskSequence documentation

* Add TaskSequence and seq_task to the __init__.py

docs/api.rst

@@ -27,6 +27,16 @@ task decorator
 
 .. autofunction:: locust.core.task
 
+TaskSequence class
+==================
+
+.. autoclass:: locust.core.TaskSequence
+	:members: locust, parent, min_wait, max_wait, wait_function, client, tasks, interrupt, schedule_task
+
+seq_task decorator
+==============
+
+.. autofunction:: locust.core.seq_task
 
 HttpSession class
 =================

docs/writing-a-locustfile.rst

@@ -255,6 +255,31 @@ its Locust instance, and the attribute :py:attr:`parent <locust.core.TaskSet.par
 parent TaskSet (it will point to the Locust instance, in the base TaskSet).
 
 
+TaskSequence class
+==================
+
+TaskSequence class is a TaskSet but its tasks will be executed in order.
+To define this order you should do the following:
+
+.. code-block:: python
+
+    class MyTaskSequence(TaskSequence):
+        @seq_task(1)
+        def first_task(self):
+            pass
+
+        @seq_task(2)
+        def second_task(self):
+            pass
+
+        @seq_task(3)
+        @task(10)
+        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.
+
 Setups, Teardowns, on_start, and on_stop
 ========================================
 

examples/browse_docs_sequence_test.py

@@ -0,0 +1,50 @@
+# This locust test script example will simulate a user
+# browsing the Locust documentation on https://docs.locust.io/
+
+import random
+from locust import HttpLocust, TaskSquence, seq_task, task
+from pyquery import PyQuery
+
+
+class BrowseDocumentationSequence(TaskSquence):
+    def on_start(self):
+        self.urls_on_current_page = self.toc_urls
+
+    # assume all users arrive at the index page
+    @seq_task(1)
+    def index_page(self):
+        r = self.client.get("/")
+        pq = PyQuery(r.content)
+        link_elements = pq(".toctree-wrapper a.internal")
+        self.toc_urls = [
+            l.attrib["href"] for l in link_elements
+        ]
+
+    @seq_task(2)
+    @task(50)
+    def load_page(self, url=None):
+        url = random.choice(self.toc_urls)
+        r = self.client.get(url)
+        pq = PyQuery(r.content)
+        link_elements = pq("a.internal")
+        self.urls_on_current_page = [
+            l.attrib["href"] for l in link_elements
+        ]
+
+    @seq_task(3)
+    @task(30)
+    def load_sub_page(self):
+        url = random.choice(self.urls_on_current_page)
+        r = self.client.get(url)
+
+
+class AwesomeUser(HttpLocust):
+    task_set = BrowseDocumentationSequence
+    host = "https://docs.locust.io/en/latest/"
+
+    # we assume someone who is browsing the Locust docs,
+    # generally has a quite long waiting time (between
+    # 20 and 600 seconds), since there's a bunch of text
+    # on each page
+    min_wait = 20 * 1000
+    max_wait = 600 * 1000

locust/__init__.py

@@ -1,4 +1,4 @@
-from .core import HttpLocust, Locust, TaskSet, task
+from .core import HttpLocust, Locust, TaskSet, TaskSequence, task, seq_task
 from .exception import InterruptTaskSet, ResponseError, RescheduleTaskImmediately
 
 __version__ = "0.8.1"

locust/core.py

@@ -58,6 +58,35 @@ 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 
@@ -419,3 +448,34 @@ def client(self):
         Locust instance.
         """
         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 it 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

locust/test/test_task_sequence_class.py

@@ -0,0 +1,74 @@
+import six
+
+from locust import InterruptTaskSet, ResponseError
+from locust.core import HttpLocust, Locust, TaskSequence, events, seq_task, task
+from locust.exception import (CatchResponseError, LocustError, RescheduleTask,
+                              RescheduleTaskImmediately)
+
+from .testcases import LocustTestCase, WebserverTestCase
+
+
+class TestTaskSet(LocustTestCase):
+    def setUp(self):
+        super(TestTaskSet, self).setUp()
+
+        class User(Locust):
+            host = "127.0.0.1"
+        self.locust = User()
+
+    def test_task_sequence_with_list(self):
+        def t1(l):
+          if l._index == 1:
+            l.t1_executed = True
+
+        def t2(l):
+          if l._index == 2:
+            l.t2_executed = True
+
+        def t3(l):
+          if l._index == 0:
+            l.t3_executed = True
+          raise InterruptTaskSet(reschedule=False)
+
+        class MyTaskSequence(TaskSequence):
+            t1_executed = False
+            t2_executed = False
+            t3_executed = False
+            tasks = [t1, t2, t3]
+
+        l = MyTaskSequence(self.locust)
+
+        self.assertRaises(RescheduleTask, lambda: l.run())
+        self.assertTrue(l.t1_executed)
+        self.assertTrue(l.t2_executed)
+        self.assertTrue(l.t3_executed)
+
+    def test_task_with_decorator(self):
+        class MyTaskSequence(TaskSequence):
+            t1_executed = 0
+            t2_executed = 0
+            t3_executed = 0
+
+            @seq_task(1)
+            def t1(self):
+              if self._index == 1:
+                self.t1_executed += 1
+
+            @seq_task(2)
+            @task(3)
+            def t2(self):
+              if self._index == 2 or self._index == 3 or self._index == 4:
+                l.t2_executed += 1
+
+            @seq_task(3)
+            def t3(self):
+              if self._index == 0:
+                self.t3_executed += 1
+              raise InterruptTaskSet(reschedule=False)
+
+        l = MyTaskSequence(self.locust)
+
+        self.assertRaises(RescheduleTask, lambda: l.run())
+        self.assertEquals(l.t1_executed, 1)
+        self.assertEquals(l.t2_executed, 3)
+        self.assertEquals(l.t3_executed, 1)