性能测试 - Locust WebSocket client

Max.Bai

2024.10

0. 背景

Locust 是性能测试工具，但是默认只支持http协议，就是默认只有http的client，需要其他协议的测试必须自己扩展对于的client，比如下面的WebSocket client。

1. WebSocket test Client

“”“
Max.Bai
Websocket Client
”“”
import json
import logging
import secrets
import threading
import time
from typing import Callable, Optional

import websocket
from locust import events

logger = logging.getLogger(__name__)


class WebSocketClient:
    def __init__(self, host: str, log_messages: bool = False):
        self._host: str = host
        self._id: str = secrets.token_hex(8)
        self._alias: Optional[str] = None
        self._ws: Optional[websocket.WebSocketApp] = None

        self.log_messages = log_messages
        self.count_recv_type = False
        self.heartbeat_auto_respond = False
        self._recv_messages: list = []
        self.messages: list = []
        self._sent_messages: list = []

    def __enter__(self):
        self.connect()
        return self

    def __exit__(self, type, value, traceback):
        self.disconnect()

    @property
    def tag(self) -> str:
        tag = f"{self._host} <{self._id}>"
        if self._alias:
            tag += f"({self._alias})"
        return tag

    def connect(
        self, alias: Optional[str] = None, headers: Optional[dict] = None, on_message: Optional[Callable] = None
    ):
        if not self._ws:
            self._alias = alias
            self._ws = websocket.WebSocketApp(
                url=self._host,
                header=headers,
                on_open=self._on_open,