Testing other systems using custom clients

Locust was built with HTTP as its main target. However, it can easily be extended to load test any request/response based system, by writing a custom client that triggers request_success and request_failure events.

Note

Any protocol libraries that you use must be gevent-friendly (use the Python socket module or some other standard library function like subprocess), or your calls will block the whole Locust process.

Some C libraries cannot be monkey patched by gevent, but allow for other workarounds. For example, if you want to use psycopg2 to performance test PostgreSQL, can use psycogreen.

Sample XML-RPC User client

Here is an example of a User class, XmlRpcUser, which provides an XML-RPC client, XmlRpcUser, and tracks all requests made:

import time
from xmlrpc.client import ServerProxy, Fault

from locust import User, task, between


class XmlRpcClient(ServerProxy):
    """
    Simple, sample XML RPC client implementation that wraps xmlrpclib.ServerProxy and
    fires locust events on request_success and request_failure, so that all requests
    gets tracked in locust's statistics.
    """

    _locust_environment = None

    def __getattr__(self, name):
        func = ServerProxy.__getattr__(self, name)

        def wrapper(*args, **kwargs):
            start_time = time.time()
            try:
                result = func(*args, **kwargs)
            except Fault as e:
                total_time = int((time.time() - start_time) * 1000)
                self._locust_environment.events.request_failure.fire(
                    request_type="xmlrpc", name=name, response_time=total_time, exception=e
                )
            else:
                total_time = int((time.time() - start_time) * 1000)
                self._locust_environment.events.request_success.fire(
                    request_type="xmlrpc", name=name, response_time=total_time, response_length=0
                )
                # In this example, I've hardcoded response_length=0. If we would want the response length to be
                # reported correctly in the statistics, we would probably need to hook in at a lower level

        return wrapper


class XmlRpcUser(User):
    """
    This is the abstract User class which should be subclassed. It provides an XML-RPC client
    that can be used to make XML-RPC requests that will be tracked in Locust's statistics.
    """

    abstract = True

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.client = XmlRpcClient(self.host)
        self.client._locust_environment = self.environment


class ApiUser(XmlRpcUser):
    host = "http://127.0.0.1:8877/"
    wait_time = between(0.1, 1)

    @task(10)
    def get_time(self):
        self.client.get_time()

    @task(5)
    def get_random_number(self):
        self.client.get_random_number(0, 100)

If you’ve written Locust tests before, you’ll recognize the class called ApiUser which is a normal User class that has a couple of tasks declared. However, the ApiUser inherits from XmlRpcUser that you can see right above ApiUser. The XmlRpcUser is marked as abstract using abstract = True which means that Locust will not try to create simulated users from that class (only of classes that extend it). XmlRpcUser provides an instance of XmlRpcClient under the client attribute.

The XmlRpcClient is a wrapper around the standard library’s xmlrpc.client.ServerProxy. It basically just proxies the function calls, but with the important addition of firing locust.event.Events.request_success and locust.event.Events.request_failure events, which will record all calls in Locust’s statistics.

Here’s an implementation of an XML-RPC server that would work as a server for the code above:

import random
import time
from xmlrpc.server import SimpleXMLRPCServer


def get_time():
    time.sleep(random.random())
    return time.time()


def get_random_number(low, high):
    time.sleep(random.random())
    return random.randint(low, high)


server = SimpleXMLRPCServer(("localhost", 8877))
print("Listening on port 8877...")
server.register_function(get_time, "get_time")
server.register_function(get_random_number, "get_random_number")
server.serve_forever()

For more examples, see locust-plugins