Locust — custom client

In this article:

- I‘ll describe a process of creating custom Locust client for different than HTTP protocol,
- I’ll show you how to create a client which uses more than one protocol.


What Locust is (very short intro):

Locust is a system for load testing written in Python (https://locust.io/). It’s easy when you need to write a few simple tests with HTTP requests only.
It starts to be fun when you want to use different protocol and write more complicated and complex tests.

I’ll assume you have a basic Python knowledge. If you don’t know Locust, please start with studying its documentation.

Let’s get started!

Basic HTTP Client:

First of all — making HTTP requests in Locust is just boring. You can find instructions and examples in documentation — https://docs.locust.io/en/latest/writing-a-locustfile.html#making-http-requests
In short: you need to create a class which inherits from HttpLocust (developed by Locust team). Then you should create your own TaskSet where you can use client via self.client instance, for example: self.client.get(’/about/’).
Locust will handle time measuring, exceptions handling, response size aggregation etc.

Cool, right?

Custom client:

What if you need to use other protocol, like FTP?
Simple — you need to implement your own custom client! Documentation: https://docs.locust.io/en/latest/testing-other-systems.html.
However in my opinion documentation is a little bit confusing and it shows a very specific example (XmlRpcClient).
To create a fully custom client you should:

  1. Create your own ProtocolClient class (of course you should replace “Protocol” with your name, like FtpClient).

I’ve created mine ProtocolClient and added __init__(self, host) method (for example to pass server configuration) in this class.
Also, in this class you should implement all other methods related to your communication as a client, for example, new_connection, upload_set, or close_connection.

2. Create time measuring wrapper (I’ll call it stopwatch).

You should add a simple function like stopwatch and a wrapper in it. This part allows you to measure the time of your own requests.

3. Add annotations above the functions’ names in the client which should be measured (i.e. upload_set).

4. Create new “Locust class” which inherits from Locust and create client object there for example

5. Now create your “User class” (ProtocolUser) and inherit from custom Locust class (ProtocolLocust) — keep it simple:

6. Implement some tasks using your custom client (i.e. self.client.upload_set())

7. TADAAAA…

So far so good.

More than one client:

Here’s the best part: what if you need to use more than 1 protocol?
In this example, we will reuse our ProtocolClient and merge it with default HTTP Client (class is named HttpSession — don’t ask me why).
In my opinion, the best way to make it is to aggregate all clients in one class. Let’s define classes with names: CustomLocust and CustomClient.
CustomLocust is rather simple:

and now the definition of CustomClient:

Small disclaimer: Take a look at the code again. Notice, that host is passed via locust parameters (— host=localhost) while starting locust. Both ports are passed via constants (which can be set based on environment variable).
If the host is different for both protocols, it would be better to use only env variables. This example shows how to use parameter passed while starting locust.

Now in our tasks, we can use clients using self.client.http… and self.client.protocol…
So you can now ask — why didn’t we define self.http and self.protocol fields in CustomLocust class?!
The answer is: because the self.client field must exist in class which inherits from Locust. We could, of course, set self.client to HTTP client and made additional self.protocol field — but the solution proposed above is way more elegant!

Example Locust files (1st with custom ProtocolClient, 2nd with both HTTP and custom ProtocolClient):

PRO tip: you can divide your code into more than one file ;)

I think this is it. Am… Yeah. Now I’m waiting for your questions, so go on and ask!