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

This comment has been minimized.

Copy link

@irshad-qb irshad-qb commented Nov 18, 2015

@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

This comment has been minimized.

Copy link
Owner Author

@yamionp yamionp commented Feb 19, 2016

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

@hirokiky

This comment has been minimized.

Copy link

@hirokiky hirokiky commented Jun 22, 2018

Big thanks for your sharing the code!

@ghost-x47

This comment has been minimized.

Copy link

@ghost-x47 ghost-x47 commented Oct 9, 2018

Thanks, this is really helpful

@Merwin89

This comment has been minimized.

Copy link

@Merwin89 Merwin89 commented Feb 11, 2020

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

@magussiro

This comment has been minimized.

Copy link

@magussiro magussiro commented Jun 19, 2020

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

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
You can’t perform that action at this time.