locust实现自定义负载形状脚本

from locust.runners import MasterRunner, WorkerRunner,STATE_STOPPING, STATE_STOPPED, STATE_CLEANUP,LocalRunner
from locust import HttpUser, TaskSet, task, events, LoadTestShape,run_single_user
from gevent._semaphore import Semaphore
import json
import traceback
import gevent
import time



'''请求成功时触发'''
def on_request_success(request_type, name, response_time, response_length):
    pass
    # print( 'success  Type: {}, Name: {}, Time: {}ms, length: {}'
    #        .format(request_type, name, response_time,response_length))

'''请求失败时触发'''
def on_request_failure(request_type, name, response_time, response_length,exception):
    print('failure  Type: {}, Name: {}, Time: {}ms, exception: {}, response_length:{}'
          .format(request_type, name, response_time,exception,response_length))

'''在执行locust类内发生异常时触发'''
def on_locust_error(locust_instance, exception, tb):
    print("error  locust_instance: {}, exception: {}, traceback: {}"
          .format(locust_instance, exception,traceback.format_tb(tb)))

#退出进程时回调
def on_quitting(environment,**kwargs):
    print("Test quit")

'''
#停止测试的时候客户端会调这个方法发送数据到主节点这边
'''
def on_test_stop(environment,**kwargs):
    if not isinstance(environment.runner, MasterRunner):
        print("Cleaning up test data")
        #节点往主节点发送的数据
        environment.runner.send_message('acknowledge_users', f"Thanks for the Cleaning up test data users!")
    else:
        users = [
            {"name": "User1"},
            {"name": "User2"},
            {"name": "User3"},
        ]
        environment.runner.send_message('test_users', users)


'''
#定义worker节点注册的消息
# Fired when the worker receives a message of type 'test_users'
'''
def setup_test_users(environment, msg, **kwargs):
    for user in msg.data:
        print(f"User {user['name']} received")
    environment.runner.send_message('acknowledge_users', f"Thanks for the {len(msg.data)} users!")


'''
#定义matser节点注册的消息
#Fired when the master receives a message of type 'acknowledge_users!!!!
'''
def on_acknowledge(msg, **kwargs):
    print("recv worker data :{}".format(msg.data))

'''
#主动退出进程时：environment.process_exit_code = 0
#判断错误率大于多少N主动退出进程
#判断响应时间大于多少N主动退出进程
#判断响应时间
'''
def checker(environment):
    while not environment.runner.state in [STATE_STOPPING, STATE_STOPPED, STATE_CLEANUP]:
        time.sleep(1)
        if environment.stats.total.fail_ratio > 0.01:
            print("Test failed due to failure ratio > 1%,code=1")
            environment.process_exit_code = 1
            '''这个语句是退出'''
            environment.runner.quit()
        elif environment.stats.total.avg_response_time > 200:
            print("Test failed due to average response time ratio > 200 ms,code=1")
            environment.process_exit_code = 1
        elif environment.stats.total.get_response_time_percentile(0.95) > 300:
            print("Test failed due to 95th percentile response time > 800 ms,code=1")
            environment.process_exit_code = 1
        else:
            environment.process_exit_code = 0
            pass
            # print("Test Normal task exit code=0")
'''
#初始化时绑定的重写方法，该类中进行了worler和master节点的消息注册
'''
def on_locust_init(environment, **kwargs):
    if not isinstance(environment.runner, MasterRunner):
        '''
        #初始化的时候注册消息，客户端的消息，往客户端的往这个类型发就行
        '''
        environment.runner.register_message('test_users', setup_test_users)
    if not isinstance(environment.runner, WorkerRunner):
        '''
        #初始化的时候注册消息，服务端的消息，往后服务端往这个类型发就行
        '''
        environment.runner.register_message('acknowledge_users', on_acknowledge)

    if isinstance(environment.runner, MasterRunner) or isinstance(environment.runner, LocalRunner):
        gevent.spawn(checker, environment)


def on_test_start(environment, **kwargs):
    '''如果运行环境不是主节点'''
    if not isinstance(environment.runner, MasterRunner):
        users = [
            {"name": "User1"},
            {"name": "User2"},
            {"name": "User3"},
        ]
        environment.runner.send_message('test_users', users)





'''''''''
#创建集合点，当locust实例产生完成时触发
'''
all_locusts_spawned = Semaphore()
#上锁
all_locusts_spawned.acquire()
'''
#生成所有locust用户时触发
'''
def on_hatch_complete(**kwargs):
    #释放锁
    all_locusts_spawned.release()

'''
#事件回调绑定
'''
events.spawning_complete.add_listener(on_hatch_complete)
events.request_success.add_listener(on_request_success)
events.request_failure.add_listener(on_request_failure)
events.user_error.add_listener(on_locust_error)
events.quitting.add_listener(on_quitting)
events.init.add_listener(on_locust_init)

'''
#主节点才触发
'''
events.test_start.add_listener(on_test_start)
events.test_stop.add_listener(on_test_stop)

header={"Content-Type": "application/json;charset=UTF-8"}
class WebsiteTasks(TaskSet):

    '''
    #进行初始化的工作，每个Locust用户开始做的第一件事
    '''
    def on_start(self):
        payload = {
            "username": "13592945579",
            "password": "5632721",
        }
        # self.client属性使用Python request库的所有方法，调用和使用方法和requests完全一致；
        with self.client.post("/user/signin",data=json.dumps(payload),catch_response=True,headers=header) as response:
            if response.status_code==200:
                response.success()
            else:
                response.failure("on_start请求失败！")


    '''
    #通过@task()装饰的方法为一个事务，方法的参数用于指定该行为的执行权重，参数越大每次被虚拟用户执行的概率越高，默认为1
    '''
    @task(1)
    def index(self):
        payload = {
            "username": "13592945579",
            "password": "5632721",
        }
        # self.client.request_name = "/blog?id=[id]" #自定义locust报告中Statistics的Name名称,进行分组显示
        with self.client.post("/user/signin",data=json.dumps(payload),catch_response=True,headers=header) as respone:#self.client属性使用Python request库的所有方法，调用和使用方法和requests完全一致；
            if respone.status_code == 200:
                respone.success()
            else:
                respone.failure("index请求失败！")


    @task(2)
    def test01(self):
        payload = {
            "username": "13592945579",
            "password": "5632721",
        }

        with self.client.post("/user/signin",data=json.dumps(payload),catch_response=True,headers=header) as respone:#self.client属性使用Python request库的所有方法，调用和使用方法和requests完全一致；
            try:
                if respone.status_code == 200:
                    respone.catch_response = True
                else:
                    respone.failure("test01请求失败！")
            except json.JSONDecodeError:
                respone.failure("Response could not be decoded as JSON!")
            except KeyError:  # 字典类型的数据，操作的变量没有该键时会报这个错
                respone.failure("Response did not contain expected key 'greeting'")


                # @task(1)
    # def about(self):
    #     self.client.get("/about/")

class WebsiteUser(HttpUser):
    # host     = "172.168.8.113:5000" #被测系统的host，在终端中启动locust时没有指定--host参数时才会用到
    host = "http://127.0.0.1:5000"
    tasks = [WebsiteTasks]          #TaskSet类，该类定义用户任务信息，必填。这里就是:WebsiteTasks类名,因为该类继承TaskSet；
    min_wait = 5000  #每个用户执行两个任务间隔时间的上下限（毫秒）,具体数值在上下限中随机取值，若不指定默认间隔时间固定为1秒
    max_wait = 15000
    # def stop(self, force=False):
    #     self.stop(True)


'''
#该类用于MyCustomShape的切片概念的时间控制，因为MyCustomShape继承了LoadTestShape
#LoadTestShape提供了tick方法，在tick方法中return (目标用户数，增长速率){当目标用户数小于前用户时，会进行递减操作，速率填写为正数}
#结合MyCustomShape实现控制用户数的概念，实现了在某个范围内进行增长，某个范围内进行稳定用户数，在某个范围内进行下降
#实现的原理，把运行波动的范围当作是个场景，当成一个切片，然后运行总时长里面就是由n个切片组合而成，那只要控制每个切片的判断，
然后让运行总时长%切片总长的余数去对比，当前余数处于切片哪个阶段，进行返回tick里面的（目标用户数，增长速率）
'''
class set_time(object):
    #定义切片的运行时间
    @property
    def sum_rtime(self):
        return self.sum_rtime

    @sum_rtime.setter
    def __set_sum_rtime__(self, value):
        self.sum_rtime = value

    @sum_rtime.getter
    def __get__sum_rtime__(self):
        return self.sum_rtime

    #定义切片的平均时间第一个节点
    @property
    def avg1_rtime(self):
        return self.avg1_rtime

    @avg1_rtime.setter
    def __set_avg_rtime__(self, value):
        self.avg1_rtime = value

    @avg1_rtime.getter
    def __get__avg_rtime__(self):
        return self.avg1_rtime

    # 定义切片的平均时间第二个节点
    @property
    def avg2_rtime(self):
        return self.avg2_rtime

    @avg2_rtime.setter
    def __set_avg_rtime__(self, value):
        self.avg2_rtime = value

    @avg2_rtime.getter
    def __get__avg_rtime__(self):
        return self.avg2_rtime

    #定义切片的开始时间
    @property
    def start_rime(self):
        return self.start_rime

    @start_rime.setter
    def __set_start_rtime__(self, value):
        self.start_rtime = value

    @start_rime.getter
    def __get__start_rtime__(self):
        return self.start_rtime

    #定义总的运行时间
    @property
    def run_time(self):
        return self.run_rime

    @run_time.setter
    def __set_run_time__(self, value):
        self.run_time = value

    @run_time.getter
    def __get__run_time__(self):
        return self.run_time



'''
#定义长期跑的时候，用户增长量和下滑量和停止时间,这个模式会覆盖config配置的模式
'''
class MyCustomShape(LoadTestShape):
    '''用户的生成速率最好是>=1,不然会出现不可控的情况，因为会出现速率小于1的情况，速率又不能等于0，所以都会进行+1，会出现到达用户数用的时间偏早的情况'''
    set_time.sum_rtime = 60 #切片运行总时间
    set_time.avg1_rtime=20 #切片运行第一个节点
    set_time.avg2_rtime=40 #切片运行第二个节点
    set_time.start_rime=0 #切片运行开始时间
    set_time.run_time=300 #整个测试时长
    user_count=10 #生成用户的顶峰数
    end_count=1#下降的目标用户数
    if int(user_count/(set_time.avg1_rtime-set_time.start_rime))==0:
        swpan1 = int(user_count / (set_time.avg1_rtime - set_time.start_rime))+1
        print("swpan1 切片开始时间到节点1时间的上升速率:{}".format(swpan1))
    else:
        swpan1=int(user_count/(set_time.avg1_rtime-set_time.start_rime))
        print("swpan1 切片开始时间到节点1时间的上升速率:{}".format(swpan1))
    if (int(user_count/(set_time.run_time-set_time.avg2_rtime)))==0:
        swpan2 = (int(user_count / (set_time.sum_rtime - set_time.avg2_rtime)))+1
        print("swpan2 节点2到切片总长时间的下降速率:{}".format(swpan2))
    else:
        swpan2=(int(user_count/(set_time.run_time-set_time.avg2_rtime)))
        print("swpan2 节点2到切片总长时间的下降速率:{}".format(swpan2))


    def tick(self):
        get_run_time = self.get_run_time()#程序正在运行的时间
        if get_run_time>=set_time.sum_rtime:
            run_time=int(get_run_time%set_time.sum_rtime)
        elif get_run_time<set_time.sum_rtime:
            run_time=get_run_time
        elif get_run_time>=set_time.run_time:
            return (None,None)
        if run_time>=set_time.start_rime  and run_time<=set_time.avg1_rtime:
            print("运行时间>开始时间 和 运行时间小于第一个节点时间 用户：{} 速率：{}".format(self.user_count,self.swpan1))
            return (self.user_count,self.swpan1)
        elif run_time>=set_time.avg1_rtime and run_time<=set_time.avg2_rtime:
            print("运行时间>第一节点时间 和 运行时间小于第二个节点时间 用户：{} 速率：{}".format(self.user_count,1))
            return (self.user_count,1)
        elif run_time>=set_time.avg2_rtime and run_time<set_time.sum_rtime:
            print("运行时间>第二节点时间 和 运行时间小于第二个节点时间 用户：{} 速率：{}".format(self.end_count,self.swpan2))
            return (self.end_count,self.swpan2)


if __name__ == '__main__':
    import os
    os.system("locust -f G:\\script_xlink\\locust_py\\script_WebsiteTasks.py  --host=http://172.168.8.114:5000")
    # os.system("locust -f G:\\script_xlink\\locust_py\\script_WebsiteTasks.py --master --host=http://172.168.8.113:5000")
    # os.system("locust -f G:\\script_xlink\\locust_py\\script_WebsiteTasks.py --master" )
    # os.system("locust -f G:\\script_xlink\\locust_py\\script_WebsiteTasks.py --worker")
    # os.system("locust --config=G:\script_xlink\locust_py\config\config.conf")
    # os.system("locust -f G:\\script_xlink\\locust_py\\script_WebsiteTasks.py --master --expect-workers=1 --headless" )

实现的效果如图：控制用户数在某个时间断进行增长，某个时间断进行平衡，某个时间断进行下降，实现控制速率