Integrating WebSocket Connections in Locust for Load Testing

Locust does not natively support WebSocket, so you can extend it by using the Python websockets library.

Step 1: Install dependencies – run pip install websockets to add the library.

Step 2: Write a Locust script – create a custom WebSocketClient class that handles connection, sending messages, receiving responses, and firing success or failure events, then define a WebSocketUser class that instantiates the client and registers an init listener to set the target WebSocket URL.

import time
import gevent
import websocket
from locust import User, task, events, between

class WebSocketClient:
    def __init__(self, host):
        self.host = host
        self.ws = None
    def connect(self):
        try:
            self.ws = websocket.create_connection(self.host)
            return "Connected"
        except Exception as e:
            return f"Connection failed: {str(e)}"
    def send(self, message):
        start_time = time.time()
        try:
            self.ws.send(message)
            response = self.ws.recv()
            total_time = int((time.time() - start_time) * 1000)
            events.request_success.fire(request_type="WebSocket", name="send", response_time=total_time, response_length=len(response))
            return response
        except Exception as e:
            total_time = int((time.time() - start_time) * 1000)
            events.request_failure.fire(request_type="WebSocket", name="send", response_time=total_time, exception=e)
    def disconnect(self):
        if self.ws:
            self.ws.close()

class WebSocketUser(User):
    abstract = True
    def __init__(self, *args, **kwargs):
        super(WebSocketUser, self).__init__(*args, **kwargs)
        self.client = WebSocketClient(self.host)

    @events.init.add_listener
    def on_locust_init(environment, **_kwargs):
        if isinstance(environment.user_classes[0], WebSocketUser):
            environment.host = "ws://yourserver/websocket/endpoint"

class MyWebSocketTest(WebSocketUser):
    wait_time = between(1, 5)

    @task
    def send_message(self):
        response = self.client.send("Hello, WebSocket!")
        print(f"Received: {response}")

    def on_start(self):
        """Called before a virtual user starts executing tasks"""
        result = self.client.connect()
        print(result)

    def on_stop(self):
        """Called after a virtual user finishes executing tasks"""
        self.client.disconnect()

Explanation – the WebSocketClient encapsulates connection logic, the WebSocketUser provides the Locust user interface, and MyWebSocketTest defines the actual test task.

Run the test – replace environment.host with your real WebSocket URL and start Locust with locust -f your_locustfile.py, then open http://localhost:8089 to configure users and begin the load test.

Using this approach you can simulate realistic WebSocket interactions in Locust for performance testing.