Locust自定义负载

269 阅读4分钟

在某些情况下,我们需要模拟复杂的用户访问行为,比如某个时间点出现流量尖峰或者按照定制的时间节奏增减用户数量。我们可以利用Locust中的LoadTestShape类来完全控制用户数量和生成速率。
继承自LoadTestShape的类。当Locust检测到这个类时,它就会自动使用这个自定义的负载模式。
在这个自定义类中,编写一个名为tick()的方法,它每秒会被调用一次,并返回期望的用户数和新增用户的速率(或返回None表示停止测试)。同时,还可以通过get_run_time()方法获取测试运行至今的时间。

from locust import LoadTestShape, HttpUser, task


class WebsiteUser(HttpUser):
    host = "https://www.example.com"
    @task
    def load_page(self):
        self.client.get("/some/page")


class TestShape(LoadTestShape):
    time_limit = 60  # 设置1分钟限制
    spawn_rate = 20  # 每秒新增用户速率

    def tick(self):
        run_time = self.get_run_time()  # 获取当前运行时间

        if run_time < self.time_limit:  # 如果未到达设定的1分钟限制
            user_count = round(run_time, -2)  # 将运行时间四舍五入到最近的百位数作为当前用户数
            return (user_count, self.spawn_rate)  # 返回用户数和新增速率

        return None  # 超过时间限制则停止测试

tick阶段指定不同的用户类型(如UserA、UserB),甚至可以根据实际需求重新启用默认命令行参数(如运行时间、spawn率和用户数)。

class StagesUsers(LoadTestShape):

    stages = [
        {"duration": 10, "users": 10, "spawn_rate": 10, "user_classes": [UserA]},
        {"duration": 30, "users": 50, "spawn_rate": 10, "user_classes": [UserA, UserB]},
        {"duration": 60, "users": 100, "spawn_rate": 10, "user_classes": [UserB]},
        {"duration": 120, "users": 100, "spawn_rate": 10, "user_classes": [UserA, UserB]},
    ]

    def tick(self):
        run_time = self.get_run_time()

        for stage in self.stages:
            if run_time < stage["duration"]:
                try:
                    tick_data = (stage["users"], stage["spawn_rate"], stage["user_classes"])
                except:
                    tick_data = (stage["users"], stage["spawn_rate"])
                return tick_data

        return None

测试开始的前10秒会创建10个UserA类型的用户,接下来的30秒会创建40个UserA和UserB混合类型的用户,以此类推,直到所有阶段结束。

如果使用自定义形状的同时还能继续使用命令行默认参数,只需将use_common_options属性设为True即可。这样,即便采用了自定义负载模式,也能确保locust接受和应用这些常见的运行参数。

现在我们结合 shape.pylocust_di.py 文件,以及 abc.py 中的抽象基类定义来分析每一行代码及执行原理。

  1. locust_di.py 文件中:
class TestShape(LoadTestShape):
    time_limit = 60  # 设置1分钟限制
    spawn_rate = 20  # 每秒新增用户速率

    def tick(self):
        run_time = self.get_run_time()  # 获取当前运行时间
        if run_time < self.time_limit:  # 如果未到达设定的1分钟限制
            user_count = round(run_time, -2)  # 将运行时间四舍五入到最近的百位数作为当前用户数
            return (user_count, self.spawn_rate)  # 返回用户数和新增速率
        return None  # 超过时间限制则停止测试
  • TestShape 类继承自 LoadTestShape 类,表示这是一个自定义负载模式。
  • 定义了两个类变量:time_limit 用于设置性能测试持续的时间(10分钟),spawn_rate 用于设置每秒钟增加用户的速率。
  • tick 方法覆盖了父类中的抽象方法,根据运行时间动态调整用户总数。如果运行时间小于设定的限制,则计算并返回对应数量的用户数以及新增速率;否则,返回 None 停止测试。
  1. shape.py 文件中:

LoadTestShape 类是 Locust 中自定义负载形状的基类,主要功能包括:

  • 使用元类 LoadTestShapeMeta 来确保子类遵循一定的规则,并允许子类通过元类特性实现一些特殊功能(比如 ABC 的行为)。
  • 类变量 runner 是一个可选的 Runner 实例引用,指向实际运行测试的 Runner 对象。
  • 类变量 abstract 标记该类为抽象类,需要子类进行扩展实现。
  • __init__ 方法初始化对象时记录起始时间。
  • reset_time 方法重置起始时间至0秒。
  • get_run_time 方法获取从测试开始到现在经过的时间(秒)。
  • get_current_user_count 方法从 Runner 实例获取当前活跃用户数量。
  • tick 方法是一个抽象方法,必须在子类中实现。它返回一个包含控制负载测试所需信息的元组(用户总数、新增速率以及可能的用户类型列表)。
  1. abc.py 文件提供了 Python 抽象基类(Abstract Base Classes, ABCs)的支持库,其中包含了一些用于定义抽象方法、属性等装饰器与类的方法,如 abstractmethod 等。

TestShape 类被实例化并在 Locust 中使用时,其执行过程如下:

  • 初始化 TestShape 类实例时会调用 __init__ 方法设置起始时间。
  • 在 Locust 运行过程中,定时调用 tick 方法以决定当前应加载多少用户,以及每秒添加或删除多少用户。
  • tick 方法首先获取运行时间,然后基于这个时间决定是否应该继续增加用户数(依据 time_limit)。
  • 当超过 time_limit 后,返回 None 表示结束测试。在此之前,按照运行时间确定的用户数和固定的 spawn_rate 返回控制参数。