Locust压测webscoket协议

Locust是一款易于使用的分布式负载测试工具。即使一个Locust节点也可以在一个进程中支持数千个用户的并发，主要通过gevent(协程)的方式。
Locust是完全基于Python，http请求完全是基于requests库。Locust支持http、https协议，还支持测试其他协议，websocket等，只要采用Python调用对应的库就可以了。
http/https采用requests;
websocket采用websocket;
记录自己的一次使用locust压测websocket协议的代码流程

#encoding：utf-8
from gevent.monkey import patch_all;patch_all()
import requests
from requests.adapters import HTTPAdapter
import os
import sys
from locust import User, task, events, constant
import time
import websocket
import ssl
import json
import uuid
import random

class AliveRequests(object):
    # 长连接
    # 会话池
    session_pool = {}

    def __enter__(self):
        self.session()
        return self

    def __exit__(self, *excinfo):
        pass

    @staticmethod
    def session():
        cur_pid = os.getpid()
        if cur_pid not in list(AliveRequests.session_pool.keys()):
            session = requests.session()
            session.mount('http://', HTTPAdapter(pool_maxsize=1000, pool_block=True))
            AliveRequests.session_pool[cur_pid] = session
        return AliveRequests.session_pool[cur_pid]

    @staticmethod
    def request(method=None, url=None, headers=None, files=None, data=None, params=None, auth=None, cookies=None, hooks=None, json=None, timeout=None):
        session = AliveRequests.session()
        request = requests.Request(method, url, headers, files, data, params, auth, cookies, hooks, json)
        while True:
            try:
                with session.send(session.prepare_request(request), timeout=(10, timeout)) as response:
                    return response
            except ConnectionError as err:
                err_obj = err.message if isinstance(err.message, Exception) else err
                if "BadStatusLine" in err_obj.message or "reset by peer" in err_obj.message:
                    continue
                else:
                    raise err
            except Exception as ex:
                raise ex

with open(os.path.join(os.path.join(os.getcwd(), 'json'), 'default.json'), 'rb') as rbf:
    json_data = json.load(rbf)

def user_login(host, username='admin', password='admin'):
    url = "http://%s:port/XXX/token" %host
    body = {"grant_type":"password","username":username,"password":password}
    headers = {"Content-Type": "application/json"}
    post_res = AliveRequests.request('POST', url, json=body, headers=headers).json()
    token = post_res["data"].get("accessToken")
    return token

def eventType_success(eventType, recvText, total_time):
    events.request_success.fire(request_type="[RECV]",
                                name=eventType,
                                response_time=total_time,
                                response_length=len(recvText))

class WebSocketClient(object):
    _locust_environment = None

    def __init__(self, host):
        self.host = host
        # 针对 WSS 关闭 SSL 校验警报
        self.ws = websocket.WebSocket(sslopt={"cert_reqs": ssl.CERT_NONE})

    def connect(self, burl):
        start_time = time.time()
        try:
            self.conn = self.ws.connect(url=burl)
        except websocket.WebSocketConnectionClosedException as e:
            total_time = int((time.time() - start_time) * 1000)
            events.request_failure.fire(
                request_type="[Connect]", name='Connection is already closed', response_time=total_time, exception=e)
        except websocket.WebSocketTimeoutException as e:
            total_time = int((time.time() - start_time) * 1000)
            events.request_failure.fire(
                request_type="[Connect]", name='TimeOut', response_time=total_time, exception=e)
        else:
            total_time = int((time.time() - start_time) * 1000)
            events.request_success.fire(
                request_type="[Connect]", name='WebSocket', response_time=total_time, response_length=0)
        return self.conn

    def recv(self):
        return self.ws.recv()

    def send(self, msg):
        self.ws.send(msg)

class WebsocketUser(User):
    abstract = True
    def __init__(self, *args, **kwargs):
        super(WebsocketUser, self).__init__(*args, **kwargs)
        self.client = WebSocketClient(self.host)
        self.client._locust_environment = self.environment


class ApiUser(WebsocketUser):
    min_wait = 1000
    max_wait = 3000

    @task(1)
    def pft(self):
        # wss 地址
        req_scb = False
        self.data = {}
        rand_num = random.randint(1000, 9999)
        rand_str = ''.join([chr(random.randint(97, 122)) for _ in range(8)])
        self.url = f"{self.client.host}/XX/XX/{rand_num}/{rand_str}/websocket"
        self.client.connect(self.url)

        self.echo_name = "%s-echo" % str(uuid.uuid4()).split("-")[0][:-1]
        topic_str = '{"topics":["%s"]}' % self.echo_name
        col = len(topic_str)
        url_kw = "subscribe"
        token = user_login(host=self.client.host.split("//")[1].split(":")[0])
        # 发送的订阅请求
        sendMsg = json.dumps(['SEND\naccessToken:%s\ndestination:/%s\ncontent-length:%d\n\n{"topics":["%s"]}\u0000' % (token, url_kw, col, self.echo_name)])
        u_start_time = time.time()
        self.client.send(sendMsg)
        try:
            # 消息接收计时
            if not req_scb:
                try:
                    for _ in range(2):
                        recv = self.client.recv()
                        if 'success' in recv:
                            req_scb = True
                            recv_j = json.loads(recv[2:-1])
                            print(recv_j)
                            break
                    u_total_time = int((time.time() - u_start_time) * 1000)
                    if 'success' not in recv:
                        print(recv)
                        events.request_failure.fire(request_type="[FAIL] WebSocket Connection Refuse",
                                                    name=self.echo_name,
                                                    response_time=u_total_time,
                                                    response_length=0,
                                                    exception=Exception(recv))
                    else:
                        events.request_success.fire(request_type='ws',
                                                    name=self.echo_name,
                                                    response_time=u_total_time,
                                                    response_length=0)
                except Exception as ex:
                    print(ex.__str__())
                    if not u_total_time:
                        u_total_time = -1
                    events.request_failure.fire(request_type="[FAIL] WebSocket Connection Refuse",
                                                name=self.echo_name,
                                                response_time=u_total_time,
                                                response_length=0,
                                                exception=ex)
            if req_scb:
                start_time = time.time()
                for jsn_inf in iter(json_data):
                    print(jsn_inf)
                    if self.echo_api(self.client.host.split("//")[1], **jsn_inf):
                        ret = str(self.client.recv())
                        print(ret)
                total_time = int((time.time() - start_time) * 1000)
                # 为每个推送过来的事件进行归类和独立计算性能指标
                try:
                    if ret.startswith('a') and 'data' in ret:
                        # ret_j = json.loads(ret)
                        # if type(ret_j) is dict:
                        #     eventType_s = jsonpath.jsonpath(ret_j, expr='$.data')
                        #     eventType_success(eventType_s, ret, total_time)
                            # 正常 OK 响应，或者其它心跳响应加入进来避免当作异常处理
                            events.request_success.fire(request_type='ws',
                                                        name=self.echo_name,
                                                        response_time=total_time,
                                                        response_length=0)
                    elif ret == '' or ret is None:
                        events.request_failure.fire(request_type="[FAIL] Data Receive is '' or None",
                                                    name=self.echo_name,
                                                    response_time=total_time,
                                                    response_length=0,
                                                    exception=Exception("ret is '' or None"))
                except websocket.WebSocketConnectionClosedException as ex:
                    events.request_failure.fire(request_type="[ERROR] WebSocketConnectionClosedException",
                                                name='Connection is already closed.',
                                                response_time=total_time,
                                                response_length=0,
                                                exception=ex)
                except websocket.WebSocketException as ex:
                    events.request_failure.fire(request_type="[ERROR] WebSocket Push Exception",
                                                name=self.echo_name,
                                                response_time=total_time,
                                                response_length=0,
                                                exception=ex)
        except Exception as ex:
            print(ex.__str__())

    def echo_api(self, host, **kwargs):
        echo_url = "http://{}XXXXX".format(host)
        body = {"topic": self.echo_name,
                "data": kwargs}
        print(echo_url)
        post_res = AliveRequests.request('POST', echo_url, json=body)
        if post_res.status_code == 200:
            print(post_res.text)
            return True


class MessageUser(User):
    task_set = ApiUser
    wait_time = constant(1000)


def usage():
    print("""
    幫助信息：
    1. ./locust_lx　必填项參數 ［非必填项參數］
    2. 參數选项：
                必填项：
                     Host IP［-i ip］
                     Host PORT［-p port］
                     Protocal [-c ws]

                非必填项：
                     用户名［-u 用户名］{-u admin}
                     登录密码［-a 登录密码］{-a admin}
                     帮助信息［-h］{-h 显示帮助信息后即退出程序}
    """)


if __name__ == "__main__":
    ip = sys.argv[sys.argv.index("-i") + 1] if "-i" in sys.argv else None
    port = sys.argv[sys.argv.index("-p") + 1] if "-p" in sys.argv else ""
    protocal = sys.argv[sys.argv.index("-c") + 1] if "-c" in sys.argv else ""
    account = sys.argv[sys.argv.index("-u") + 1] if "-u" in sys.argv else 'admin'
    password = sys.argv[sys.argv.index("-a") + 1] if "-a" in sys.argv else 'admin'
    if len(sys.argv) < 2 or "-h" in sys.argv or not ip:
        usage()
        exit(0)
    os.system("locust -f locust_ws.py -H {}".format("{}://{}:{}".format(protocal, ip, port)))