Locust 壓測websocket協議

三只松鼠發表於2024-07-24
Web協議

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')

相關文章