Testing websockets

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)