locust-WebSocket压测

连接WS的库有的是支持异步IO的，项目中我们推荐这样的库，但是压测时还是要选择同步的库。

异步

安装

pip install websockets

代码示例

import asyncio
import websockets
import json
import random

async def mytest():
    async with websockets.connect('wss://sockettest.xhkjedu.com/ws') as websocket:
        num = random.randint(0, 10000000)
        msg = {
            "b": {"num": num},
            "c": 123456,
        }

        msgstr = json.dumps(msg)

        await websocket.send(msgstr)
        print(f"↑: {msgstr}")

        greeting = await websocket.recv()
        print(f"↓: {greeting}")

asyncio.get_event_loop().run_until_complete(mytest())

同步

官网地址

https://pypi.org/project/websocket-client/

安装

pip install websocket-client

示例

from websocket import create_connection
import json
import random

ws = create_connection("wss://sockettest.xhkjedu.com/ws")

num = random.randint(0, 10000000)
msg = {
    "b": {"num": num},
    "c": 123456,
}

msgstr = json.dumps(msg)
print("Sending " + msgstr)
ws.send(msgstr)

result = ws.recv()
print("Received '%s'" % result)
ws.close()

import websocket

def on_message(ws, message):
    print(ws)
    print(message)

def on_error(ws, error):
    print(ws)
    print(error)

def on_close(ws):
    print(ws)
    print("### closed ###")

websocket.enableTrace(True)
ws = websocket.WebSocketApp("ws://127.0.0.1:8888/track",
                            on_message=on_message,
                            on_error=on_error,
                            on_close=on_close)

ws.run_forever()

Websocket压测

Jmeter要测试websocket接口，需要先下载安装一个websocket samplers by peter doornbosch的插件

而locust因为是代码实现，所以可以进行任何的测试，引用相应的库即可。

from locust import User, task, events
import time
from websocket import create_connection
import json
import random

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

def fail_call(name, total_time, e):
    events.request_failure.fire(
        request_type="[Fail]",
        name=name,
        response_time=total_time,
        response_length=0,
        exception=e,
    )

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

    def connect(self, burl):
        self.ws = create_connection(burl)

    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):
    host = "wss://sockettest.xhkjedu.com/"

    @task(1)
    def pft(self):
        # wss 地址
        self.url = 'wss://sockettest.xhkjedu.com/ws'
        print("连接前")
        start_time = time.time()
        try:
            self.client.connect(self.url)
            print("连接后")
            # 发送的订阅请求
            num = random.randint(0, 10000000)
            msg = {
                "b": {"num": num},
                "c": 123456,
            }

            msgstr = json.dumps(msg)

            self.client.send(msgstr)
            print(f"↑: {msgstr}")

            greeting = self.client.recv()
            print(f"↓: {greeting}")

        except Exception as e:
            total_time = int((time.time() - start_time) * 1000)
            fail_call("Send", total_time, e)
        else:
            total_time = int((time.time() - start_time) * 1000)
            success_call("Send", "success", total_time)