测试 - 基于 Locust 的 WebSocket-server 压力测试

一、Locust 的基本学习

Locust 官方文档

1-0 最简单的事例

from locust import Locust, TaskSet, task

class MyTaskSet(TaskSet):
    @task
    def my_task(self):
        print("executing my_task")

class MyLocust(Locust):
    task_set = MyTaskSet
    min_wait = 5000
    max_wait = 15000
    
# locust -f locust_file.py MyLocust

1-1 基本使用

from locust import HttpLocust, TaskSet, task, Locust
from locust.clients import HttpSession
from locust.exception import LocustError


class UserBehavior(TaskSet):
    def on_start(self):
        """ on_start is called when a Locust start before any task is scheduled """
        self.login()

    def on_stop(self):
        """ on_stop is called when the TaskSet is stopping """
        self.logout()

    def login(self):
        self.client.post("/login", {"username": "ellen_key", "password": "education"})

    def logout(self):
        self.client.post("/logout", {"username": "ellen_key", "password": "education"})

    @task(2)
    def index(self):
        self.client.get("/")

    @task(1)
    def profile(self):
        self.client.get("/profile")


# Then we have a HttpLocust class which represents a user,
# where we define how long a simulated user should wait between executing tasks,
# as well as what TaskSet class should define the user’s “behaviour”.
class WebsiteUser(HttpLocust):
    # Locust - task_set.run() - UserBehavior.run() - TaskSet.run()
    task_set = UserBehavior
    # 执行事物之间用户等待时间的下界，单位毫秒
    min_wait = 5000
    max_wait = 9000


'''
RUN WAY : locust -f locust_files/my_locust_file.py --host=http://example.com

locust -f supperdiancan.py --no-web -c 10 -r 3 --run-time 10s
-f   locust_file.py文件(locust任务文件)
-c   指定要生成的Locust用户数
-r   每秒生成的用户数
-n   用于0.8用于-n指定请求数
--run-time 或-t   指定测试的运行时间
'''

1-2 基于内置 http 请求的源码分析

• 基本流程分析
  • WebsiteUser - Locust 子类，用于配制相关信息；可自定义名字，若但文件存在多个该类，则可使用命令进行指定执行。
  • UserBehavior - TaskSet 子类，用于配制相关任务，每个user执行一次。
  • HttpLocust - Locust 子类，用于初始化 client 交与 WebsiteUser
  • HttpSession - 使用 requests 模块，实现基于 https 协议的连接

二、实现基于 WebSocket 协议的压力测试

• 仿写思路
  • 实现 ws 连接类，类似于 HttpSession
  • 继承 Locust 类，实现 client 的覆盖，类比HttpLocust
  • UserBehavior 内实现服务链接，一个用户一个链接
  • WebsiteUser 配置信息，将 task_set 指向 UserBehavior

import gzip
import json
import time

import websocket
from locust import TaskSet, events, Locust, task


class WebSocketClient(object):

    def connect(self, url):
        try:
            self.ws = websocket.WebSocketApp(url)
            self.ws.on_message = self.on_message
            self.ws.on_error = self.on_error
            self.ws.on_close = self.on_close
        except websocket.WebSocketTimeoutException as e:
            events.request_failure.fire(request_type="web_socket", name='ws', response_time=time.time(), exception=e)
        else:
            events.request_success.fire(request_type="web_socket", name='ws', response_time=time.time(),
                                        response_length=0)
        return self.ws

    def on_message(self, message):
        if isinstance(message, bytes):
            message = gzip.decompress(message).decode("utf-8")
        message = json.loads(message)
        print(message)
        if isinstance(message, dict):
            if 'ping' in message:
                pong = {'pong': message['ping']}
                self.ws.send(json.dumps(pong))

    def on_error(self, error):
        print('!!! error !!!', error)

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

    def on_open(self):
        print('opened')
        data = {"sub": {'category': 0, 'star': []}}
        self.ws.send(json.dumps(data))


class WebSocketLocust(Locust):
    def __init__(self):
        super(WebSocketLocust, self).__init__()
        self.client = WebSocketClient()


class UserBehavior(TaskSet):
    def on_start(self):
        print('--------- task start ------------')
        self.url = 'ws://localhost:8000/v1/ws/marketpair'

    def on_stop(self):
        print('---------- task stop ------------')

    @task
    def test_ws(self):
        ws = self.client.connect(self.url)
        ws.run_forever()


class WebsiteUser(WebSocketLocust):
    task_set = UserBehavior
    min_wait = 5000
    max_wait = 9000

# locust -f market_pair_connect.py --host=http://127.0.0.1/

三、可视化界面测试

默认测试页面 ：http://127.0.0.1:8089/