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() |
This comment has been minimized.
This comment has been minimized.
@irshad-qb You are right. But result screens become luxurious a little. |
This comment has been minimized.
This comment has been minimized.
Big thanks for your sharing the code! |
This comment has been minimized.
This comment has been minimized.
Thanks, this is really helpful |
This comment has been minimized.
This comment has been minimized.
How do i send in cookies and messages? Please do reply asap...stuck in a scenario |
This comment has been minimized.
This comment has been minimized.
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
This comment has been minimized.
@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.