Skip to content

Instantly share code, notes, and snippets.

Embed
What would you like to do?
Websocket Locust Sample. locustfile and Echo/Chat Server
# -*- coding:utf-8 -*-
from __future__ import absolute_import
from __future__ import unicode_literals
from __future__ import print_function
import json
import uuid
import time
import gevent
from websocket import create_connection
import six
from locust import HttpLocust, TaskSet, task
from locust.events import request_success
class EchoTaskSet(TaskSet):
def on_start(self):
self.user_id = six.text_type(uuid.uuid4())
ws = create_connection('ws://127.0.0.1:5000/echo')
self.ws = ws
def _receive():
while True:
res = ws.recv()
data = json.loads(res)
end_at = time.time()
response_time = int((end_at - data['start_at']) * 1000000)
request_success.fire(
request_type='WebSocket Recv',
name='test/ws/echo',
response_time=response_time,
response_length=len(res),
)
gevent.spawn(_receive)
def on_quit(self):
self.ws.close()
@task
def sent(self):
start_at = time.time()
body = json.dumps({'message': 'hello, world', 'user_id': self.user_id, 'start_at': start_at})
self.ws.send(body)
request_success.fire(
request_type='WebSocket Sent',
name='test/ws/echo',
response_time=int((time.time() - start_at) * 1000000),
response_length=len(body),
)
class ChatTaskSet(TaskSet):
def on_start(self):
self.user_id = six.text_type(uuid.uuid4())
ws = create_connection('ws://127.0.0.1:5000/chat')
self.ws = ws
def _receive():
while True:
res = ws.recv()
data = json.loads(res)
end_at = time.time()
response_time = int((end_at - data['start_at']) * 1000000)
request_success.fire(
request_type='WebSocket Recv',
name='test/ws/chat',
response_time=response_time,
response_length=len(res),
)
gevent.spawn(_receive)
def on_quit(self):
self.ws.close()
@task
def sent(self):
start_at = time.time()
body = json.dumps({'message': 'hello, world', 'user_id': self.user_id, 'start_at': start_at})
self.ws.send(body)
request_success.fire(
request_type='WebSocket Sent',
name='test/ws/chat',
response_time=int((time.time() - start_at) * 1000000),
response_length=len(body),
)
class EchoLocust(HttpLocust):
task_set = EchoTaskSet
min_wait = 0
max_wait = 100
class ChatLocust(HttpLocust):
task_set = ChatTaskSet
min_wait = 0
max_wait = 100
Flask==0.10.1
gevent==1.0.2
gevent-websocket==0.9.4
greenlet==0.4.7
itsdangerous==0.24
Jinja2==2.7.3
linecache2==1.0.0
locustio==0.7.2
MarkupSafe==0.23
mock==1.0.1
msgpack-python==0.4.6
pyzmq==14.6.0
requests==2.7.0
six==1.9.0
traceback2==1.4.0
unittest2==1.0.1
websocket-client==0.31.0
Werkzeug==0.10.4
# -*- coding: utf-8 -*-
from __future__ import absolute_import
from __future__ import unicode_literals
from __future__ import print_function
from collections import defaultdict
import json
from geventwebsocket.handler import WebSocketHandler
from gevent.pywsgi import WSGIServer
from flask import Flask, request
from werkzeug.exceptions import abort
app = Flask(__name__)
ctr = defaultdict(int)
@app.route('/echo')
def echo():
ws = request.environ['wsgi.websocket']
if not ws:
abort(400)
while True:
message = ws.receive()
if message is not None:
r = json.loads(message)
ctr[r['user_id']] += 1
ws.send(message)
@app.route('/report')
def report():
return '\n'.join(['{}:\t{}'.format(user_id, count) for user_id, count in ctr.items()])
socket_handlers = set()
@app.route('/chat')
def chat():
ws = request.environ['wsgi.websocket']
socket_handlers.add(ws)
while True:
message = ws.receive()
for socket_handler in socket_handlers:
try:
socket_handler.send(message)
except:
socket_handlers.remove(socket_handler)
if __name__ == '__main__':
http_server = WSGIServer(('', 5000), app, handler_class=WebSocketHandler)
http_server.serve_forever()
@irshad-qb
Copy link

@yamionp : Why you are calling request_success.fire after sending the message ? It only require on receiving right ? I am new to locust, So please don’t mind if this question has no relevance.

@yamionp
Copy link
Author

yamionp commented Feb 19, 2016

@irshad-qb You are right. But result screens become luxurious a little.

@hirokiky
Copy link

Big thanks for your sharing the code!

@ghost-x47
Copy link

Thanks, this is really helpful

@Merwin89
Copy link

How do i send in cookies and messages? Please do reply asap...stuck in a scenario

@magussiro
Copy link

hi yamionp , i'm new to locust and python, i'm using your script to achieve 1000 over chat room stress test, it's work, but when I shutdown locust it's always show some greenlet error, would you tell me how to fix it? it would be very helpful to me ,very thanks

@jeffotoni
Copy link

and the client so we can test our server, is there any example out there?

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment