Locust 压测websocket协议

Locust 自带的HttpLocust类只支持http，并不支持websocket，但实现websocket压测只需要三步。

1、第一步需要自己写一个websocket client的类，实现连接、发送、接受、关闭连接等方法

import websocket

class WebSocketClient(object):
    def __init__(self, host):
        self.host = host
        self.ws = websocket.WebSocket()

    def connect(self, url):
        self.ws.connect(url=self.host + url)

        return self.ws

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

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

    def close(self):
        self.ws.close()

2、继承 User 类（当前最新Locust版本是2.29.1，如果是老版本，继承的是 Locust 类）

class WebsocketLocust(User):
    # 定义抽象类
    abstract = True
    def __init__(self, *args, **kwargs):
        super(WebsocketLocust, self).__init__(*args, **kwargs)
        # 输入真实的ws接口ip和端口
        self.client = WebSocketClient('ws://ip:port')

3、继承 WebsocketLocust 类，初始化User，定义测试行为

# 定义用户行为
class UserTask(WebsocketLocust):
    def on_start(self):
        '''初始化数据，每个虚拟用户只执行一次'''
        self.url = "/v1/test"
        self.ws = self.client.connect(self.url)
        self.params ={}

    # 多个接口用@task(1)、@task(2)分配权限，按1:2流量分配     
    @task
    def test(self):
        start_time = time.time()
        try:
            self.ws.send(json.dumps(self.params))
            r = self.ws.recv()
            total_time = time.time() - start_time
            result = json.loads(r)['data']
            assert result[3] != '未查询到相关信息'
            events.request.fire(request_type="websockt", name=self.url, response_time=total_time * 1000,
                                response_length=len(r))

        except AssertionError:
            events.request.fire(request_type="websockt", name=self.url, response_time=total_time * 1000,
                                response_length=0, exception=f"断言错误,response:{result}")

        except Exception as e:
            total_time = time.time() - start_time
            events.request.fire(request_type="websockt", name=self.url, response_time=total_time*1000, exception=e)

    def on_stop(self):
        '''销毁数据，每个虚拟用户只执行一次'''
        self.ws.close()

总结起来就是：

from locust import task, between,User,events
import os, json, time,websocket

class WebSocketClient(object):
    def __init__(self, host):
        self.host = host
        self.ws = websocket.WebSocket()

    def connect(self, url):
        self.ws.connect(url=self.host + url)

        return self.ws

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

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

    def close(self):
        self.ws.close()
  
class WebsocketLocust(User):
    # 定义抽象类
    abstract = True
    def __init__(self, *args, **kwargs):
        super(WebsocketLocust, self).__init__(*args, **kwargs)
        # 输入真实的ws接口ip和端口
        self.client = WebSocketClient('ws://ip:port')
        
# 定义用户行为
class UserTask(WebsocketLocust):
    def on_start(self):
        '''初始化数据，每个虚拟用户只执行一次'''
        self.url = "/v1/test"
        self.ws = self.client.connect(self.url)
        self.params ={}

    # 多个接口用@task(1)、@task(2)分配权限，按1:2流量分配     
    @task
    def test(self):
        start_time = time.time()
        try:
            self.ws.send(json.dumps(self.params))
            r = self.ws.recv()
            total_time = time.time() - start_time
            result = json.loads(r)['data']
            assert result[3] != '未查询到相关信息'
            events.request.fire(request_type="websockt", name=self.url, response_time=total_time * 1000,
                                response_length=len(r))

        except AssertionError:
            events.request.fire(request_type="websockt", name=self.url, response_time=total_time * 1000,
                                response_length=0, exception=f"断言错误,response:{result}")

        except Exception as e:
            total_time = time.time() - start_time
            events.request.fire(request_type="websockt", name=self.url, response_time=total_time*1000, exception=e)

    def on_stop(self):
        '''销毁数据，每个虚拟用户只执行一次'''
        self.ws.close()
       
   
if __name__ == "__main__":
    os.system('locust -f stress_test.py --web-port 8090')