Locust 之@task和@tag装饰器梳理

        在Locust中的User类和TaskSet类中，定义一个task都是使用@task装饰器来定义一个任务，下面我们来看看它是如何实现的。

一、task装饰器

        首先来看看task装饰器的源码

def task(weight: Union[TaskT, int] = 1) -> Union[TaskT, Callable[[TaskT], TaskT]]:
    """
    Used as a convenience decorator to be able to declare tasks for a User or a TaskSet
    inline in the class. Example::
        class ForumPage(TaskSet):
            @task(100)
            def read_thread(self):
                pass
            @task(7)
            def create_thread(self):
                pass
            @task(25)
            class ForumThread(TaskSet):
                @task
                def get_author(self):
                    pass
                @task
                def get_created(self):
                    pass
    """
 
    def decorator_func(func):
        if func.__name__ in ["on_stop", "on_start"]:
            logging.warning(
                "You have tagged your on_stop/start function with @task. This will make the method get called both as a task AND on stop/start."
            )  # this is usually not what the user intended
        if func.__name__ == "run":
            raise Exception(
                "User.run() is a method used internally by Locust, and you must not override it or register it as a task"
            )
        func.locust_task_weight = weight
        return func
 
    """
    Check if task was used without parentheses (not called), like this::
        @task
        def my_task()
            pass
    """
    if callable(weight):
        func = weight
        weight = 1
        return decorator_func(func)
    else:
        return decorator_func

        task函数接受TaskT或int类型对象作为参数，该参数代表task的权重。返回类型为TaskT或一个可调用对象数组。

        1、首先判断如果weight参数是一个可调用对象，则先把weight存入一个临时变量func中，然后设置weight=1。接着调用task函数内部定义的一个函数decorator_func，传入参数为func。调用decorator_func函数返回的是func函数本身，并且在func函数上动态增加一个locust_task_weight属性，属性值为weight;

        总结一下：当weight是一个可调用对象时，@task做的仅仅是给传入task的可调用对象weight动态增加了一个值为1的locust_task_weight属性；

from typing import (
    TYPE_CHECKING,
    Callable,
    List,
    Union,
    TypeVar,
    Optional,
    Type,
    overload,
    Dict,
    Set,
)
 
TaskT = TypeVar("TaskT", Callable[..., None], Type["TaskSet"])
 
 
def task(weight: Union[TaskT, int] = 1) -> Union[TaskT, Callable[[TaskT], TaskT]]:
 
    def decorator_func(func):
        func.locust_task_weight = weight
        return func
 
    if callable(weight):
        func = weight
        weight = 1
        return decorator_func(func)
    else:
        return decorator_func
 
 
@task
def myTask(str):
    print(str)
 
 
if __name__ == '__main__':
    myTask("hello")
    print(myTask.locust_task_weight)

其中

@task
def myTask(str):
    print(str)

 本质上等于myTask = task(myTask)

task(myTask)函数返回了decorator_func(myTask)，而decorator_func(myTask)又返回了myTask，然后把最后的myTask赋值给一个同名的函数变量myTask。

        接着调用myTask()，其实还是调用myTask()本身。

      2、如果weight参数类型为int型，则直接返回task函数内部的decorator_func函数对象并赋值给myTask。

        当调用myTask(xxx)时，其实就是调用decorator_func(xxx)，其函数内部会在xxx对象上增加locust_task_weight属性，最后返回xxx。

        在locaust中，这种情况其实就是给user活TaskSet对象上增加了locust_task_weight属性，如下所示

# 首页
class UserBehavior(TaskSet):
 
    @task(1)
    def test_02(self):
        """
        查看首页商品详情
        :return:
        """
        param = {}
        with self.client.get('/', json=param) as response:
            print("查看首页商品详情")

        综上，task装饰器函数本质上就是给其装饰的函数或对象上增加locust_task_weight属性。

如果@task直接装饰函数A，则A的locust_task_weight=1， 如果@task参数为int类型，则参数self的locust_task_weight=weight。

二、tag装饰器

def tag(*tags: str) -> Callable[[TaskT], TaskT]:
    """
    Decorator for tagging tasks and TaskSets with the given tag name. You can
    then limit the test to only execute tasks that are tagged with any of the
    tags provided by the :code:`--tags` command-line argument. Example::
        class ForumPage(TaskSet):
            @tag('thread')
            @task(100)
            def read_thread(self):
                pass
            @tag('thread')
            @tag('post')
            @task(7)
            def create_thread(self):
                pass
            @tag('post')
            @task(11)
            def comment(self):
                pass
    """
 
    def decorator_func(decorated):
        if hasattr(decorated, "tasks"):
            decorated.tasks = list(map(tag(*tags), decorated.tasks))
        else:
            if "locust_tag_set" not in decorated.__dict__:
                decorated.locust_tag_set = set()
            decorated.locust_tag_set |= set(tags)
        return decorated
 
    if len(tags) == 0 or callable(tags[0]):
        raise ValueError("No tag name was supplied")
 
    return decorator_func

可见tag装饰器期望传入一个字符串类型参数。然后返回其内部的decorator_func函数对象。

decorator_func = tag("xxx"); mytask = decorator_func; 算了太乱了，被搞混了以后再分析吧。

三、何处使用locust_task_weight呢？

def get_tasks_from_base_classes(bases, class_dict):
    """
    Function used by both TaskSetMeta and UserMeta for collecting all declared tasks
    on the TaskSet/User class and all its base classes
    """
    new_tasks = []
    for base in bases:
        if hasattr(base, "tasks") and base.tasks:
            new_tasks += base.tasks
 
    if "tasks" in class_dict and class_dict["tasks"] is not None:
        tasks = class_dict["tasks"]
        if isinstance(tasks, dict):
            tasks = tasks.items()
 
        for task in tasks:
            if isinstance(task, tuple):
                task, count = task
                for _ in range(count):
                    new_tasks.append(task)
            else:
                new_tasks.append(task)
 
    for item in class_dict.values():
        if "locust_task_weight" in dir(item):
            for i in range(item.locust_task_weight):
                new_tasks.append(item)
 
    return new_tasks

        可以看到，在这个地方会检查item中是否有locust_task_weight属性，如果有，则根据其权重值使用了一个for循环来不断地往task队列中加入对应权重值的任务。