Merge pull request #1358 from Trouv/mark-tasks

Add task marking for running more specific tests

docs/api.rst

@@ -27,6 +27,11 @@ task decorator
 
 .. autofunction:: locust.task
 
+tag decorator
+==============
+
+.. autofunction:: locust.tag
+
 SequentialTaskSet class
 =======================
 

docs/changelog.rst

@@ -45,6 +45,25 @@ attribute instead:
         tasks = [MyTaskSet]
 
 
+Task tagging
+------------
+
+A new :ref:`tag feature <tagging-tasks>` has been added that makes it possible to include/exclude tasks during 
+a test run.
+
+Tasks can be tagged using the :py:func:`@tag <locust.tag>` decorator:
+
+.. code-block:: python
+
+    class WebUser(User):
+        @task
+        @tag("tag1", "tag2")
+        def my_task(self):
+            ...
+
+And tasks can then be specified/excluded using the ``--tags``/``-T`` and ``--exclude-tags``/``-E`` command line arguments. 
+
+
 Environment variables changed
 -----------------------------
 

docs/writing-a-locustfile.rst

@@ -193,6 +193,51 @@ that looks like this::
 and then Python's ``random.choice()`` is used pick tasks from the list.
 
 
+.. _tagging-tasks:
+
+Tagging tasks
+-------------
+
+By tagging tasks using the `tag <locust.tag>` decorator, you can be picky about what tasks are
+executed during the test using the :code:`--tags` and :code:`--exclude-tags` arguments.  Consider
+the following example:
+
+.. code-block:: python
+
+    from locust import User, constant, task, tag
+
+    class MyUser(User):
+        wait_time = constant(1)
+
+        @tag('tag1')
+        @task
+        def task1(self):
+            pass
+
+        @tag('tag1', 'tag2')
+        @task
+        def task2(self):
+            pass
+
+        @tag('tag3')
+        @task
+        def task3(self):
+            pass
+
+        @task
+        def task4(self):
+            pass
+
+If you started this test with :code:`--tags tag1`, only *task1* and *task2* would be executed
+during the test. If you started it with :code:`--tags tag2 tag3`, only *task2* and *task3* would be
+executed.
+
+:code:`--exclude-tags` will behave in the exact opposite way. So, if you start the test with
+:code:`--exclude-tags tag3`, only *task1*, *task2*, and *task4* will be executed. Exclusion always
+wins over inclusion, so if a task has a tag you've included and a tag you've excluded, it will not
+be executed.
+
+
 
 TaskSet class
 =============
@@ -237,9 +282,9 @@ A TaskSet can also be inlined directly under a User/TaskSet class using the @tas
         class MyTaskSet(TaskSet):
             ...
 
-
 The tasks of a TaskSet class can be other TaskSet classes, allowing them to be nested any number 
 of levels. This allows us to define a behaviour that simulates users in a more realistic way. 
+
 For example we could define TaskSets with the following structure::
 
     - Main user behaviour
@@ -312,11 +357,19 @@ User instance.
 Referencing the User instance, or the parent TaskSet instance
 ---------------------------------------------------------------
 
-A TaskSet instance will have the attribute :py:attr:`locust <locust.TaskSet.user>` point to 
+A TaskSet instance will have the attribute :py:attr:`user <locust.TaskSet.user>` point to
 its User instance, and the attribute :py:attr:`parent <locust.TaskSet.parent>` point to its
 parent TaskSet instance.
 
 
+Tags and TaskSets
+------------------
+You can tag TaskSets using the `tag <locust.tag>` decorator in a similar way to normal tasks, as
+described `above <tagging-tasks>`, but there are some nuances worth mentioning. Tagging a TaskSet
+will automatically apply the tag(s) to all of the TaskSet's tasks. Furthermore, if you tag a task
+within a nested TaskSet, locust will execute that task even if the TaskSet isn't tagged.
+
+
 .. _sequential-taskset:
 
 SequentialTaskSet class

locust/__init__.py

@@ -5,7 +5,7 @@
 
 from .user.sequential_taskset import SequentialTaskSet
 from .user import wait_time
-from .user.task import task, TaskSet
+from .user.task import task, tag, TaskSet
 from .user.users import HttpUser, User
 from .user.wait_time import between, constant, constant_pacing
 

locust/argument_parser.py

@@ -267,6 +267,22 @@ def setup_parser_arguments(parser):
         env_var="LOCUST_MASTER_NODE_PORT",
     )
 
+    tag_group = parser.add_argument_group("Tag options", "Locust tasks can be tagged using the @tag decorator. These options let specify which tasks to include or exclude during a test.")
+    tag_group.add_argument(
+        '-T', '--tags',
+        nargs='*',
+        metavar='TAG',
+        env_var="LOCUST_TAGS",
+        help="List of tags to include in the test, so only tasks with any matching tags will be executed"
+    )
+    tag_group.add_argument(
+        '-E', '--exclude-tags',
+        nargs='*',
+        metavar='TAG',
+        env_var="LOCUST_EXCLUDE_TAGS",
+        help="List of tags to exclude from the test, so only tasks with no matching tags will be executed"
+    )
+
     stats_group = parser.add_argument_group("Request statistics options")
     stats_group.add_argument(
         '--csv',
@@ -401,4 +417,4 @@ def get_parser(default_config_files=DEFAULT_CONFIG_FILES):
 
 
 def parse_options(args=None):
-    return get_parser().parse_args(args=args)
+    return get_parser().parse_args(args=args)

locust/env.py

@@ -3,6 +3,7 @@
 from .stats import RequestStats
 from .runners import LocalRunner, MasterRunner, WorkerRunner
 from .web import WebUI
+from .user.task import filter_tasks_by_tags
 
 
 class Environment:
@@ -15,6 +16,12 @@ class Environment:
     user_classes = []
     """User classes that the runner will run"""
 
+    tags = None
+    """If set, only tasks that are tagged by tags in this list will be executed"""
+
+    exclude_tags = None
+    """If set, only tasks that aren't tagged by tags in this list will be executed"""
+
     stats = None
     """Reference to RequestStats instance"""
 
@@ -51,6 +58,8 @@ class Environment:
     def  __init__(
         self, *,
         user_classes=[],
+        tags=None,
+        exclude_tags=None,
         events=None, 
         host=None, 
         reset_stats=False, 
@@ -65,13 +74,17 @@ def  __init__(
             self.events = Events()
 
         self.user_classes = user_classes
+        self.tags = tags
+        self.exclude_tags = exclude_tags
         self.stats = RequestStats()
         self.host = host
         self.reset_stats = reset_stats
         self.step_load = step_load
         self.stop_timeout = stop_timeout
         self.catch_exceptions = catch_exceptions
         self.parsed_options = parsed_options
+
+        self._filter_tasks_by_tags()
 
     def _create_runner(self, runner_class, *args, **kwargs):
         if self.runner is not None:
@@ -114,7 +127,7 @@ def create_worker_runner(self, master_host, master_port):
             master_host=master_host,
             master_port=master_port,
         )
-    
+
     def create_web_ui(self, host="", port=8089, auth_credentials=None, tls_cert=None, tls_key=None):
         """
         Creates a :class:`WebUI <locust.web.WebUI>` instance for this Environment and start running the web server
@@ -130,3 +143,16 @@ def create_web_ui(self, host="", port=8089, auth_credentials=None, tls_cert=None
         """
         self.web_ui = WebUI(self, host, port, auth_credentials=auth_credentials, tls_cert=tls_cert, tls_key=tls_key)
         return self.web_ui
+
+    def _filter_tasks_by_tags(self):
+        """
+        Filter the tasks on all the user_classes recursively, according to the tags and
+        exclude_tags attributes
+        """
+        if self.tags is not None:
+            self.tags = set(self.tags)
+        if self.exclude_tags is not None:
+            self.exclude_tags = set(self.exclude_tags)
+
+        for user_class in self.user_classes:
+            filter_tasks_by_tags(user_class, self.tags, self.exclude_tags)

locust/main.py

@@ -93,6 +93,8 @@ def create_environment(user_classes, options, events=None):
     """
     return Environment(
         user_classes=user_classes,
+        tags=options.tags,
+        exclude_tags=options.exclude_tags,
         events=events,
         host=options.host,
         reset_stats=options.reset_stats,

locust/test/test_locust_class.py

@@ -3,7 +3,7 @@
 from gevent.pool import Group
 
 from locust.exception import InterruptTaskSet, ResponseError
-from locust import HttpUser, User, TaskSet, task, between, constant
+from locust import HttpUser, User, TaskSet, task, constant
 from locust.env import Environment
 from locust.exception import (CatchResponseError, LocustError, RescheduleTask,
                               RescheduleTaskImmediately, StopUser)

locust/test/test_runners.py

@@ -65,6 +65,8 @@ def __init__(self):
         self.hatch_rate = 5
         self.num_users = 5
         self.host = '/'
+        self.tags = None
+        self.exclude_tags = None
         self.master_host = 'localhost'
         self.master_port = 5557
         self.master_bind_host = '*'