locustfile.py

import time
import random
import string
import websocket
from locust import HttpUser, TaskSet, task, between, User


class WebSocketClient:
    def __init__(self, host):
        self.host = host
        self.ws = None

    def connect(self):
        self.ws = websocket.create_connection(self.host)

    def send(self, message):
        self.ws.send(message)

    def receive(self):
        return self.ws.recv()

    def close(self):
        if self.ws:
            self.ws.close()


class WebSocketTaskSet(TaskSet):
    def __init__(self, parent):
        super().__init__(parent)
        self.ws_client = None

    def on_start(self):
        # Generate random session_id
        length = 12
        session_id = ''.join(random.choices(string.ascii_letters + string.digits, k=length))

        # Connect
        self.ws_client = WebSocketClient(f"ws://127.0.0.1:1865/ws/{session_id}")
        self.ws_client.connect()

    def on_stop(self):
        self.ws_client.close()

    @task
    def send_message(self):
        start_time = time.time()
        try:
            self.ws_client.send('{"text": "Hello"}')

            # Read until we get a response without "chat_token"
            while True:
                response = self.ws_client.receive()
                if "chat_token" not in response:
                    break


            # Decide success or failure
            if "chat" in response:
                # Mark as success: pass `exception=None`
                self.user.environment.events.request.fire(
                    request_type="websocket",
                    name="send_message",
                    response_time=(time.time() - start_time) * 1000,  # ms
                    response_length=len(response),
                    exception=None
                )
            else:
                # Mark as failure: pass an Exception
                self.user.environment.events.request.fire(
                    request_type="websocket",
                    name="send_message",
                    response_time=(time.time() - start_time) * 1000,
                    response_length=len(response),
                    exception=Exception("No 'chat' in response")
                )
        except Exception as e:
            # If anything goes wrong, also mark as failure
            self.user.environment.events.request.fire(
                request_type="websocket",
                name="send_message",
                response_time=(time.time() - start_time) * 1000,
                response_length=0,
                exception=e
            )


class MyWebSocketUser(HttpUser):
    tasks = [WebSocketTaskSet]
    wait_time = between(5, 10)