Distributed load testing: Extending locus.io for connection oriented services (2/5)

I have to admit that when talking about testing scripts I'm a python fan because of the simplicity, low overhead and availability of libraries (Ruby could be a good option too).  So, the first think I did when started to build my distributed load testing environment was searching google for "python distributed load testing" and as usual google didn't dissapoint me and that was the beginning of fun with locust!

locust is python testing framework where you write your tests in python and the execution can be scheduled in one or multiple machines.    It has facilities for HTTP testing but can be extended to other services and provides a web interface to get easy access to the progress and results of the tests.

After digging for a while in the documentation (not that good) and the source code (nice code) I decided to give it a try and create my first non HTTP test.   When creating a Locust class (that class defines the configuration of the tests) you can specify a client to override the default HTTP client:

class MyLocust(Locust):
    task_set = MyTaskSet
    min_wait = 500
    max_wait = 500

    def __init__(self):
        super(MyLocust, self).__init__()

        self.client = MyProtocolLocustClient(host=self.host)

class MyProtocolLocustClient(object):
    def __init__(self, host):
        self.host = host
        self.connection = MyProtocolConnection(host)

    def ping(self):
        request_meta = {}
        request_meta["start_time"] = time.time()
        request_meta["method"] = 'MESSAGE'
        request_meta["name"] = 'PING'

        self.connection.send_message(self.connection.id(), "HI")

        try:
            response = self.connection.recv_message()

            if payload != "HI":
                raise Exception()

            request_meta["response_time"] = (time.time() - request_meta["start_time"]) * 1000

            events.request_success.fire(
                    method=request_meta["method"],
                    name=request_meta["name"],
                    response_time=request_meta["response_time"],
                    response_length=0
                )
        except Exception as e:
            events.request_failure.fire(
                    method=request_meta["method"],
                    name=request_meta["name"],
                    response_time=0,
                    exception=e,
                    response=None,
                )


That client object is available to all the locust TaskSets that you will define later with your specific tests.  For example:

class MyTaskSet(TaskSet):
    @task
    def my_task(l):
        l.client.ping()

And that's all, just put that code in a locustfile.py file and run it from the command line:
 locust -f locustfile.py -H xxxx.yyyy.com

And open a browser in http://localhost:8089 to start the test and get the results and you should get something like this: