← Load Test Directory

Testing Websockets

How to use LoadForge to load test websockets based applications

You are now browsing the LoadForge locust test directory. You can use these tests as a starting point for your own tests, or use our AI wizard to generate one automatically.

World

Overview

Websocket testing is naturally more complicated than simple webserver, or web application load testing. However, with LoadForge and locust files you can connect to and load test websockets easily.

For an easier implementation, see our SocketIO websockets test guide.

First example

from locust import HttpUser, TaskSet, task, between,events
from websocket import create_connection
import time
import random

from uuid import uuid4
class UserBehavior(TaskSet):
    statements = [
        'Hello',
        'How are you?',
        'This chat is nice',
    ]

    def on_start(self):
        uri = "wss://YOUR_WS_URL_HERE:443/chat/user" + str(uuid4().hex)
        self.ws = create_connection(uri)
        self.ws.send(str(uuid4().hex))

    def on_quit(self):
        self.ws.disconnect()

    @task(1)
    def say_hello(self):
        start_at = time.time()

        body = {"message": "Hello", "customData": {"language": "en"}}

        self.ws.send(str(body))
        events.request_success.fire(
            request_type='WebSocket Sent',
            name='test/ws/echo',
            response_time=int((time.time() - start_at) * 1000000),
            response_length=len(body),
        )

    @task(2)
    def say_random(self):
        start_at = time.time()

        statement = random.choice(self.statements)
        body = {"message": statement, "customData": {"language": "en"}}

        self.ws.send(str(body))
        events.request_success.fire(
            request_type='WebSocket Sent',
            name='test/ws/echo',
            response_time=int((time.time() - start_at) * 1000000),
            response_length=len(body),
        )

class WebsiteUser(HttpUser):
    tasks = [UserBehavior]
    wait_time = between(5, 15)

Second example

import time
import json
import gevent

from uuid import uuid4

from locust import HttpUser, TaskSet, task, events, User
import websocket

class SocketClient(object):
    def __init__(self, host):
        self.host = host
        self.session_id = uuid4().hex
        self.connect()

    def connect(self):
        self.ws = websocket.WebSocket()
        self.ws.settimeout(10)
        self.ws.connect(self.host)

        events.quitting += self.on_close

        self.attach_session()

    def attach_session(self):
        payload = {'action': 'attach_session', 'session_id': self.session_id}
        json_data = self.send_with_response(payload)

        assert json_data['action'].lower() == 'attach_session'
        assert json_data['message_status'].lower() == 'ok'
        assert json_data['session_id'] == self.session_id

    def send_with_response(self, payload):
        json_data = json.dumps(payload)

        g = gevent.spawn(self.ws.send, json_data)
        g.get(block=True, timeout=2)
        g = gevent.spawn(self.ws.recv)
        result = g.get(block=True, timeout=10)

        json_data = json.loads(result)
        return json_data

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

    def send(self, payload):
        message_id = uuid4().hex
        payload.update({'_messageId': message_id,
                        'session_id': self.session_id})
        start_time = time.time()
        e = None
        try:
            data = self.send_with_response(payload)
            assert data['_messageId'] == message_id
            assert data['session_id'] == self.session_id
        except AssertionError as exp:
            e = exp
        except Exception as exp:
            e = exp
            self.ws.close()
            self.connect()
        elapsed = int((time.time() - start_time) * 1000)
        if e:
            events.request_failure.fire(request_type='sockjs', name='send',
                                        response_time=elapsed, exception=e)
        else:
            events.request_success.fire(request_type='sockjs', name='send',
                                        response_time=elapsed,
                                        response_length=0)

class WSBehavior(TaskSet):
    @task(1)
    def action(self):
        data = {
            "action": "do_stuff",
            "param": "123",
        }
        self.client.send(data)

class WSUser(User):
    task_set = WSBehavior
    min_wait = 1000
    max_wait = 3000

    def __init__(self, *args, **kwargs):
        super(WSUser, self).__init__(*args, **kwargs)
        self.client = SocketClient('ws://%s/rt/websocket' % self.host)

Locust Test Example

LoadForge is powered by locust, meaning open source locust users can copy this script as well. If you are a locust user, consider importing your script to LoadForge to supercharge your testing!

Ready to run your test?
Start your first test within minutes.