locust mixed rest and websocket performance testing on the same locustfile

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

The problem may be that you're putting multiple Users in your SequentialTaskSet tasks. I haven't seen anyone do that before so I don't know what the behavior would be. I'd recommend making your MySocketUserClass a generic class object that doesn't inherit from User. You could make it take the environment on init so it could still fire events. Then instead of having WebSocketCall be yet another class, just make the task a function inside the SequentialTaskSet. You can actually do the same for RestCallUser. Then you make a FastHttpUser with the SequentialTaskSet as its tasks. See this SO question for an example:

https://stackoverflow.com/a/63573818/7355637

Alternatively, you can also just not use SequentialTaskSet and instead make the sent1 task more generic and call your websocket my_task stuff first inside the same task. A single task can include multiple function calls inside it. Locust will still fire whatever events appropriately.

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.