Integrating Multiple Custom Clients in Locust for HTTP and WebSocket Testing

In Locust you can integrate multiple custom clients to simulate different request types or services within a single test script, allowing combined performance testing of HTTP and WebSocket endpoints.

Implementation steps : define a client class for each service, instantiate these clients in a user class, and configure Locust run parameters such as user count and spawn rate.

Example code demonstrates a WebSocketClient class, an HttpClient class, and a MultiProtocolUser class that uses both clients. The user class defines separate tasks for HTTP and WebSocket requests and includes on_start and on_stop hooks to manage connections.

import websocket
import requests
from locust import User, task, between, events

# WebSocket client
class WebSocketClient:
    def __init__(self, host):
        self.host = host
        self.ws = None
    def connect(self):
        self.ws = websocket.create_connection(self.host)
        return "Connected"
    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()

# HTTP client (illustrative, Locust already provides HttpUser)
class HttpClient:
    def __init__(self, base_url):
        self.base_url = base_url
    def get(self, path):
        start_time = time.time()
        try:
            response = requests.get(f"{self.base_url}{path}")
            total_time = int((time.time() - start_time) * 1000)
            if response.status_code == 200:
                events.request_success.fire(request_type="HTTP", name=path, response_time=total_time, response_length=len(response.text))
            else:
                events.request_failure.fire(request_type="HTTP", name=path, response_time=total_time, exception=Exception(f"Unexpected status code {response.status_code}"))
            return response.text
        except Exception as e:
            total_time = int((time.time() - start_time) * 1000)
            events.request_failure.fire(request_type="HTTP", name=path, response_time=total_time, exception=e)

class MultiProtocolUser(User):
    abstract = True
    def __init__(self, *args, **kwargs):
        super(MultiProtocolUser, self).__init__(*args, **kwargs)
        self.http_client = HttpClient(self.host)  # HTTP base URL
        self.ws_client = WebSocketClient("ws://your_websocket_server")
    @task(1)
    def http_task(self):
        self.http_client.get("/some_path")
    @task(2)
    def ws_task(self):
        self.ws_client.send("Hello from WebSocket client")
    def on_start(self):
        """Called when a virtual user starts"""
        print("Connecting to WebSocket server")
        self.ws_client.connect()
    def on_stop(self):
        """Called when a virtual user stops"""
        print("Disconnecting from WebSocket server")
        self.ws_client.disconnect()

Run Locust with a command such as locust -f your_locustfile.py --headless -u 100 -r 10 -t 1m to start 100 users, spawn 10 users per second, and run for one minute.

Notes : ensure proper resource cleanup (e.g., close WebSocket connections in on_stop ), add robust exception handling for each client, and extend the pattern by defining additional client classes for new services.