locust mixed rest and websocket performance testing on the same locustfile

Question

I'm trying to write a test with the following requirementes:

1. the test has two calls, the calls has to be exectuted in sequence
2. the first call is a REST call (http post)
3. the second one is a websocket call

How I can implement this in the same locustfile? Is this possible?

This is my test:

import time
import json
import logging
import re
import gevent
import websocket

from locust import User, FastHttpUser, SequentialTaskSet, task, between

headers = [
    "auth: headers",
]


class MySocketUserClass(User):
    def connect(self, host: str, header=[], **kwargs):
        self.ws = websocket.create_connection(host, header=header, **kwargs)
        gevent.spawn(self.receive_loop)

    def send(self, body, context={}):
        self.environment.events.request.fire(
            request_type="WSS",
            name='test',
            response_time=None,
            response_length=len(body),
            exception=None,
            context={**self.context(), **context},
        )
        logging.debug(f"WSS: {body}")
        self.ws.send(json.dumps(body))

    def on_message(self, message):  # override this method in your subclass for custom handling
        print(message)
        self.environment.events.request.fire(
            request_type="WSR",
            name='test',
            response_time=0,
            response_length=len(message),
            exception=None,
            context=self.context(),
        )

    def receive_loop(self):
        while True:
            message = self.ws.recv()
            logging.debug(f"WSR: {message}")
            self.on_message(message)
    
    def sleep_with_heartbeat(self, seconds):
        while seconds >= 0:
            gevent.sleep(min(15, seconds))
            seconds -= 15
            self.send({})

class WebSocketCall(MySocketUserClass):
    @task
    def my_task(self):

        self.connect("wss://example.com", header=headers)

        # example of subscribe
        self.send({'example': 'payload'})

        # wait for additional pushes, while occasionally sending heartbeats, like a real client would
        self.sleep_with_heartbeat(10)

    def on_message(self, message):
        # TO BE IMPLEMENTED
        pass

class RestCall(FastHttpUser):
    default_headers = {
        "accept": "application/json",
        "accept-encoding": "gzip, deflate, br",
        "accept-language": "en-US,en;q=0.9,it;q=0.8",
        "content-type": "application/json",
        "sec-ch-ua": '"Chromium";v="110", "Not A(Brand";v="24", "Microsoft Edge";v="110"',
        "sec-ch-ua-mobile": "?0",
        "sec-ch-ua-platform": '"Windows"',
        "sec-fetch-dest": "empty",
        "sec-fetch-mode": "cors",
        "sec-fetch-site": "same-origin",
        "user-agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/110.0.0.0 Safari/537.36 Edg/110.0.1587.50",
        "origin": "https://example.com",
        "referer": "https://example.com",
    }

    def call_rest(self):
        self.iam = 'rest'

        with self.rest(
            "POST",
            "https://example.com",
            json={
                "example": "payload"
            },
        ) as resp:
            # CHECK RESPONSE
            print(resp)

    @task
    def sent1(self):
        self.call_rest()

class UserBehaviour(SequentialTaskSet):
    tasks = [RestCall, WebSocketCall]

but only the RestCall is executed

Answer 1

You can use a single User class for this.

As SocketIOUser’s documentation states, it is:

A User that includes a socket io websocket connection. You could easily use this a template for plain WebSockets, socket.io just happens to be my use case. You can use multiple inheritance to combine this with an HttpUser (class MyUser(HttpUser, SocketIOUser)

So something like:

class MySocketIOUser(HttpUser, SocketIOUser):
    @task
    def my_task(self):
        self.connect("wss://...")
        self.client.post(...)
        self.send("something")

See https://github.com/SvenskaSpel/locust-plugins/blob/master/examples/socketio_ex.py for more about the WebSocket-parts.