Locust1.4 websocket压力测试

locust 1.4后用user 代替了locust

测试http使用的 HttpUser是依赖requests包，基类为User

测试Websocket 需要借助于 websocket 继承User类，

import json
import logging
import os
import queue
import random
import string
import sys
import traceback
import websocket
import time

from locust import events, TaskSet, task, User, constant_pacing
from mix_loadtest.log import log_config

log = log_config(f_level=logging.INFO, c_level=logging.INFO, out_path='../../../logs',
                 filename='tst_wb_log2', fix=True)[0]

class WebSocketClient(object):
    
    def __init__(self, host):
        self.host = host
        self.ws = websocket.WebSocket()
        self.name = "wsTest"
    
    def record_result(self, response_time, response_length=0, exception_msg=None):
        self.name = "wsTest"
        if exception_msg:
            events.request_failure.fire(request_type="WS", name=self.name, response_time=response_time, exception=exception_msg,
                                        response_length=len(exception_msg))
        else:
            events.request_success.fire(request_type="WS", name=self.name, response_time=response_time, response_length=response_length)
    
    def connect(self, burl, request_name='urlweb'):
        self.name = request_name
        start_time = time.time()
        try:
            self.conn = self.ws.connect(url=burl)
        except websocket.WebSocketTimeoutException as e:
            total_time = int((time.time() - start_time) * 1000)
            self.record_result(response_time=total_time, exception_msg=e)
        else:
            total_time = int((time.time() - start_time) * 1000)
            self.record_result(response_time=total_time)
        log.info(f'{request_name} connect success ....')
        return self.conn
    
    def recv(self):
        return self.ws.recv()
    
    def send(self, msg):
        self.ws.send(msg)
    
    def rec_msg(self, expect_str=None, time_out=500, forever=False, time_out_per=60, run_user=None):
        pass


class SupperSC(TaskSet):
    
    def login_in(self, host, user, password):
        login_url = 'http://localhost:/login'   #登录URL
        body = {
            "user": user,
            "password": password,
        }
        import requests
        res = requests.post(login_url, json=body)
        return res.json().get('token')
    
    def on_start(self):
        data = self.user.queueData.get()  # 获取队列里的数据
        self.username, password = data.get('username'), data.get('password')
        token = self.login_in(self.user.host, self.username, password)
    
    @task(1)
    def test_wb(self):
        host = self.client.host
        self.url = 'ws://{}/websocket'.format(host)
        self.client.connect(self.url, self.username)
        # wb数据send/rec 及逻辑
        
class WSUser(User):
    host = '192.168.1.1:10219'  #待测主机
    wait_time = constant_pacing(1)  # 单个用户执行间隔时间
    tasks = [SupperSC]

    queueData = queue.Queue()  # 队列实例化
    for count in range(3000):  # 循环数据生成
        data = {
            "username": f'tst_user_{count}',
            "password": "Admin@123",
        }
        queueData.put_nowait(data)
    def __init__(self, *args, **kwargs):
        super(WSUser, self).__init__(*args, **kwargs)
        self.client = WebSocketClient(self.host)

if __name__ == "__main__":
    os.system("locust -f l_runner_wb2.py")