Websocket client for locust.io (SockJS)

locustfile.py

# Run with 
# $ locust -H ws-server:port -P 8080 -f locustfile.py
# then open the browser on localhost:8080

# requirements.txt
# locust==0.0
# locustio==0.14.5
# websocket==0.2.1

import time
import json
import gevent
from uuid import uuid4

from locust import HttpLocust, TaskSet, task, between, ResponseError, events, Locust
import websocket

class SocketClient(object):
    def __init__(self, host):
        self.host = host
        self.session_id = uuid4().hex
        self.connect()

    def connect(self):
        self.ws = websocket.WebSocket()
        self.ws.settimeout(10)
        self.ws.connect(self.host)
        events.quitting += self.on_close

    def send_with_response(self, payload):
        #print('sending data {}'.format(payload))
        json_data = json.dumps(payload)

        g = gevent.spawn(self.ws.send, json_data)
        g.get(block=True, timeout=2)
        g = gevent.spawn(self.ws.recv)
        result = g.get(block=True, timeout=10)

        #print('received data {}'.format(result))
        json_data = json.loads(result)
        return json_data

    def on_close(self):
        self.ws.close()

    def send(self, payload):
        start_time = time.time()
        e = None
        try:
            data = self.send_with_response(payload)
        except AssertionError as exp:
            e = exp
        except Exception as exp:
            e = exp
            self.ws.close()
            # why send a connect here?
            self.connect()
        elapsed = int((time.time() - start_time) * 1000)
        if e:
            events.request_failure.fire(request_type='sockjs', name='send',
                                        response_time=elapsed, response_length=0, exception=e)
        else:
            events.request_success.fire(request_type='sockjs', name='send',
                                        response_time=elapsed,
                                        response_length=0)

class WSBehavior(TaskSet):
    @task(1)
    def action(self):
        data = {
            "id": 123,
            "val": 456,
        }
        self.client.send(data)


class WSUser(HttpLocust):
    task_set = WSBehavior
    wait_time = between(5.0, 9.0)

    def __init__(self, *args, **kwargs):
        super(WSUser, self).__init__(*args, **kwargs)
        self.client = SocketClient('ws://%s/chat' % self.host)