Skip to content

Instantly share code, notes, and snippets.

@thstarshine
Forked from yamionp/locustfile.py
Created June 30, 2020 02:44
Show Gist options
  • Star 0 You must be signed in to star a gist
  • Fork 0 You must be signed in to fork a gist
  • Save thstarshine/60c39ec3854c730fdb0ecf115c2e2722 to your computer and use it in GitHub Desktop.
Save thstarshine/60c39ec3854c730fdb0ecf115c2e2722 to your computer and use it in GitHub Desktop.
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
certifi==2019.3.9
chardet==3.0.4
Click==7.0
Flask==1.0.2
gevent==1.4.0
gevent-websocket==0.10.1
greenlet==0.4.15
idna==2.8
itsdangerous==1.1.0
Jinja2==2.10.1
linecache2==1.0.0
locustio==0.11.0
MarkupSafe==1.1.1
mock==3.0.5
msgpack==0.6.1
msgpack-python==0.5.6
pyzmq==18.0.1
requests==2.21.0
six==1.12.0
traceback2==1.4.0
unittest2==1.1.0
urllib3==1.24.2
websocket-client==0.57.0
Werkzeug==0.15.1
# -*- 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()
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment