性能测试 - 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 格式的信息处理

相关推荐
金丝猴也是猿3 分钟前
如何在Flutter上实现高性能的动态模板渲染
websocket·网络协议·tcp/ip·http·网络安全·https·udp
swift开发pk OC开发6 分钟前
Flutter性能优化细节
websocket·网络协议·tcp/ip·http·网络安全·https·udp
网络抓包与爬虫7 分钟前
Flutter 环境搭建、常用指令、开发细节
websocket·网络协议·tcp/ip·http·网络安全·https·udp
天天进步201530 分钟前
Python项目-基于Flask的个人博客系统设计与实现(1)
开发语言·python·flask
安然无虞32 分钟前
31天Python入门——第20天:魔法方法详解
开发语言·后端·爬虫·python
靠近彗星43 分钟前
基于 Vue + Django + MySQL 实现个人博客/CMS系统
前端·vue.js·python·mysql·django
励志成为大佬的小杨1 小时前
pytorch模型的进阶训练和性能优化
人工智能·pytorch·python
m0_490240672 小时前
软件自动化测试(1):python+selenium自动化测试环境搭建
开发语言·python·selenium
橘猫云计算机设计2 小时前
基于ssm的食物营养成分数据分析平台设计与实现(源码+lw+部署文档+讲解),源码可白嫖!
后端·python·信息可视化·数据挖掘·数据分析·django·毕业设计
色的归属感3 小时前
Flutter完整开发实战详解(三、 打包与填坑篇)
websocket·网络协议·tcp/ip·http·网络安全·https·udp