Locust 压测websocket协议

  Locust 自带的HttpLocust类只支持http，并不支持websocket，但实现websocket压测只需要三步，配合events事件的使用，让压测报告数据输出准确。

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):
        self.ws.send(json.dumps(self.params))
        r = self.ws.recv()
        result = json.loads(r)['data']
        assert result[3] != '未查询到相关信息'

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

 以上三步就完成了 Locust 支持 websocket 的压测。如果要自定义 断言、异常等统计，需要用到 events 事件。

 

4、events 事件

  Locust 2.15版本已经移除了 events 事件中的 request_success、request_failure 功能，新版统一用 request 替代，异常事件用 user_error 。

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),response=result)

        except AssertionError:
            events.request.fire(request_type="websockt", name=self.url, response_time=total_time * 1000,
                                response_length=0, exception=f"断言错误,response:{result}")
        
        except OSError as e:
            events.user_error.fire(user_instance=UserTask, exception=e, tb=sys.exc_info()[2])
        
        except socket.timeout:
            events.user_error.fire(user_instance=UserTask, exception='Timeout', tb=sys.exc_info()[2])

        except Exception as e:
            events.user_error.fire(user_instance=UserTask, exception=e, tb=sys.exc_info()[2])

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

 

5、打印日志

log = logging.getLogger('locust')

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()
            
            # 使用日志
            log.info('this is a test')
            log.error('this is a test')
            
            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),response=result)

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

        except OSError as e:
            events.user_error.fire(user_instance=UserTask, exception=e, tb=sys.exc_info()[2])            
            
        except socket.timeout:
            events.user_error.fire(user_instance=UserTask, exception='Timeout', tb=sys.exc_info()[2])

        except Exception as e:
            events.user_error.fire(user_instance=UserTask, exception=e, tb=sys.exc_info()[2])

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

if __name__ == "__main__":
    os.system('locust -f stress_test_ws.py --web-port 8090')
    # 不输出控制台，只输出到log文件
    # os.system('locust -f stress_test_ws.py --web-port 8090 --logfile=locust.log')        

# 输出：
'''
[2024-08-20 11:49:58,204] admin-pc/INFO/locust: this is a test
[2024-08-20 11:49:58,204] admin-pc/ERROR/locust: this is a test
'''

 

把上面的代码总结起来就是：

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

log = logging.getLogger('locust')

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
    # 每个虚拟用户执行完一次任务后随机1~5秒等待，再执行下一个任务
    # wait_time = between(1, 5)
    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()
            log.info(r)
            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),response=result)

        except AssertionError:
            events.request.fire(request_type="websockt", name=self.url, response_time=total_time * 1000,
                                response_length=0, exception=f"断言错误,response:{result}")
        
        except OSError as e:
            events.user_error.fire(user_instance=UserTask, exception=e, tb=sys.exc_info()[2])
            
        except socket.timeout:
            events.user_error.fire(user_instance=UserTask, exception='Timeout', tb=sys.exc_info()[2])

        except Exception as e:
            events.user_error.fire(user_instance=UserTask, exception=e, tb=sys.exc_info()[2])

    def on_stop(self):
        '''销毁数据，每个虚拟用户只执行一次'''
        self.ws.close()
       
   
if __name__ == "__main__":
    os.system('locust -f stress_test.py --web-port 8090')
    # 只输出log日志，不输出控制台
    # os.system('locust -f stress_test_ws.py --web-port 8090 --logfile=locust.log')

 

压测报告：