性能测试 - Locust WebSocket client

Max.Bai

2024.10

0. 背景

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

1. WebSocket test Client
python 复制代码
"""
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,
                on_message=on_message if on_message else self._on_message,
                on_close=self._on_close,
            )
            thread = threading.Thread(target=self._ws.run_forever)
            thread.daemon = True
            thread.start()
            time.sleep(3)
        else:
            logger.warning("An active WebSocket connection is already established.")

    def is_connected(self) -> bool:
        return self._ws is not None

    def disconnect(self):
        if self._ws:
            self._ws.close()
            self._alias = None
        else:
            logger.warning("No active WebSocket connection established.")

    def _on_open(self, ws):
        logger.debug(f"[WebSocket] {self.tag} connected.")
        events.request.fire(
            request_type="ws_client",
            name="connect",
            response_time=0,
            response_length=0,
        )

    def _on_message(self, ws, message):
        recv_time = time.time()
        recv_time_ms = int(recv_time * 1000)
        recv_time_ns = int(recv_time * 1000000)
        logger.debug(f"[WebSocket] {self.tag} message received: {message}")
        if self.log_messages:
            self._recv_messages.append(message)
        self.messages.append(message)

        # public/respond-heartbeat
        if self.heartbeat_auto_respond:
            if "public/heartbeat" in message:
                self.send(message.replace("public/heartbeat", "public/respond-heartbeat"))

        if self.count_recv_type:
            try:
                msg = json.loads(message)
                id = str(msg.get("id", 0))
                if len(id) == 13:
                    resp_time = recv_time_ms - int(id)
                elif len(id) == 16:
                    resp_time = (recv_time_ns - int(id)) / 1000
                elif len(id) > 13:
                    resp_time = recv_time_ms - int(id[:13])
                else:
                    resp_time = 0
                method = msg.get("method", "unknown")
                code = msg.get("code", "unknown")
                error = msg.get("message", "unknown")
                # send_time = int(msg.get("nonce", 0))
                if method in ["public/heartbeat", "private/set-cancel-on-disconnect"]:
                    events.request.fire(
                        request_type="ws_client",
                        name=f"recv {method}",
                        response_time=0,
                        response_length=len(msg),
                    )
                elif code == 0:
                    events.request.fire(
                        request_type="ws_client",
                        name=f"recv {method} {code}",
                        # response_time=recv_time - send_time,
                        response_time=resp_time,
                        response_length=len(msg),
                    )
                else:
                    events.request.fire(
                        request_type="ws_client",
                        name=f"recv {method} {code}",
                        response_time=resp_time,
                        response_length=len(msg),
                        exception=error,
                    )
            except Exception as e:
                events.request.fire(
                    request_type="ws_client",
                    name="recv error",
                    response_time=0,
                    response_length=len(msg),
                    exception=str(e),
                )

    def _on_close(self, ws, close_status_code, close_msg):
        logger.debug(f"[WebSocket] {self.tag} closed.")
        self._ws = None

        events.request.fire(
            request_type="ws_client",
            name="close",
            response_time=0,
            response_length=0,
        )

    def set_on_message(self, on_message: Callable):
        self._ws.on_message = on_message

    def send(self, message: str):
        if self._ws:
            self._ws.send(data=message)
            if self.log_messages:
                self._sent_messages.append(message)
            logger.debug(f"[WebSocket] {self.tag} message sent: {message}")
        else:
            logger.warning(f"No active [WebSocket] {self.tag} connection established.")
            raise ConnectionError("No active [WebSocket] connection established.")

    def clear(self):
        self._recv_messages = []
        self._sent_messages = []
        self.messages = []

    def expect_messages(
        self,
        matcher: Callable[..., bool],
        count: int = 1,
        timeout: int = 10,
        interval: int = 1,
    ) -> list:
        """Expect to receive one or more filtered messages.

        Args:
            matcher (Callable): A matcher function used to filter the received messages.
            count (int, optional): Number of messages to be expected before timeout. Defaults to 1.
            timeout (int, optional): Timeout in seconds. Defaults to 10.
            interval (int, optional): Interval in seconds. Defaults to 1.

        Returns:
            list: A list of messages filtered by the matcher.
        """

        deadline: float = time.time() + timeout
        result: list = []  # messages filtered by the matcher
        seen: list = []  # messages already seen by the matcher to be excluded from further matching

        while time.time() < deadline:
            snapshot: list = [*self._recv_messages]

            for element in seen:
                if element in snapshot:
                    snapshot.remove(element)

            result.extend(filter(matcher, snapshot))
            if len(result) >= count:
                break

            seen.extend(snapshot)
            time.sleep(interval)

        if len(result) < count:
            logger.warning(
                f"({self.tag}) Expected to receive {count} messages, but received only {len(result)} messages."
            )

        return result
2. 如何使用
python 复制代码
class PrivateWsUser(User):

    def on_start(self):
        self.ws_client=WebSocketClient("wss://abc.pp.com/chat", log_message=True)
        self.ws_client.connect()


    @task
    def send_hello()
        self.ws_client.send("hello world")
    
    
            
3. 扩展

可自行扩展on_message 方法,上面的on_message 方法是json 格式的信息处理

相关推荐
闲人编程1 分钟前
OpenCV图像轮廓分析完全指南
python·opencv·图像识别
神仙别闹9 分钟前
基于Python+Neo4j实现新冠信息挖掘系统
开发语言·python·neo4j
navyDagger10 分钟前
GAN生成对抗网络数学原理解释并实现MNIST数据集生产(附代码演示)
人工智能·python
没有梦想的咸鱼185-1037-166314 分钟前
【降尺度】ChatGPT+DeepSeek+python+CMIP6数据分析与可视化、降尺度技术与气候变化的区域影响、极端气候分析
python·chatgpt·数据分析
berryyan17 分钟前
傻瓜教程安装Trae IDE用AI撰写第一个AKShare接口脚本
python·trae
灏瀚星空34 分钟前
从基础到实战的量化交易全流程学习:1.3 数学与统计学基础——概率与统计基础 | 基础概念
笔记·python·学习·金融·概率论
Hellohistory39 分钟前
HOTP 算法与实现解析
后端·python
伊织code41 分钟前
cached-property - 类属性缓存装饰器
python·缓存·cache·装饰器·ttl·property·cached-property
明明跟你说过1 小时前
深度学习常见框架:TensorFlow 与 PyTorch 简介与对比
人工智能·pytorch·python·深度学习·自然语言处理·tensorflow
搏博1 小时前
专家系统的基本概念解析——基于《人工智能原理与方法》的深度拓展
人工智能·python·深度学习·算法·机器学习·概率论