基于locust的websocket压测
背景：
locust默认内部只封装httplocust；使用的是requests中的session进行了封装；如果我想测试其它协议怎么办，比如websocket , grpc；我们只要重写一个实例给client即可：
重写WebSocketClient类（主要用来替换掉self.client的http实例）
class WebSocketClient(object):
def __init__(self, host):
self.host = host
self.ws = websocket.WebSocket()
def connect(self, burl):
start_time = time.time()
try:
self.conn = self.ws.connect(url=burl)
except websocket.WebSocketTimeoutException as e:
total_time = int((time.time() - start_time) * 1000)
events.request_failure.fire(request_type="websockt", name='urlweb', response_time=total_time, exception=e)
else:
total_time = int((time.time() - start_time) * 1000)
events.request_success.fire(request_type="websockt", name='urlweb', response_time=total_time, response_length=0)
return self.conn
def recv(self):
return self.ws.recv()
def send(self, msg):
self.ws.send(msg)
注意：该类中定义了，websocket的常用操作，链接、接收、发送；最主要是events.request_failure.fire和events.request_success.fire这两个用来收集性能数据，如果不写报告收集不到性能数据
2、重写一个HttpLocust类，我们这里叫做WebsoketLoscust类
class WebsocketLocust(Locust):
def __init__(self, *args, **kwargs):
super(WebsocketLocust, self).__init__(*args, **kwargs)
self.client = WebSocketClient(self.host)
注意：WebsocketLocust从Locust继承； 这里主要是将self.client重新实例成，我们第一部写好的websocketClient实例
3、编写TaskSet类
class SupperDianCan(TaskSet):
@task
def test_baidu(self):
self.url = 'wss://xxxxxx.xxxx.com/cart/chat?sid=11303&table_no=103&user_id=ofZjWs40HxEzvV08l6m4PnqGbxqc_2_1_&version=2'
self.data = {}
self.client.connect(self.url)
while True:
recv = self.client.recv()
print(recv)
if eval(recv)['type'] == 'keepalive':
self.client.send(recv)
else:
self.client.send(self.data)
注意：此类就是任务类，跟http的写法一样，只是这里用的self.client.xxxx已经变成了我们自已重写的websocket类，将原来的requests http替换了
4/编写站点类
class WebsiteUser(WebsocketLocust):
task_set = SupperDianCan
min_wait=5000
max_wait=9000
注意：站点类从第二步中的locust继承完整代码1：
from locust import Locust, events, task, TaskSet
import websocket
import time
import gzip
class WebSocketClient():
def __init__(self, host):
self.host = host
#self.port = port
class WebSocketLocust(Locust):
def __init__(self, *args, **kwargs):
self.client = WebSocketClient("1xx.xx.xx.85")
class UserBehavior(TaskSet):
@task(1)
def buy(self):
try:
ws = websocket.WebSocket()
# self.ws.connect("ws://xx:8807")
ws.connect("ws://xxxx.com/r1/xx/ws")
start_time = time.time()
#self.ws.send('{"url":"/buy","data":{"id":"123","issue":"20170822","doubled_num":2}}')
#result = self.ws.recv()
send_info = '{"sub": "market.ethusdt.kline.1min","id": "id10"}'
# send_info = '{"event":"subscribe", "channel":"btc_usdt.deep"}'
while True:
# time.sleep(5)
# ws.send(json.dumps(send_info))
ws.send(send_info)
while (1):
compressData = ws.recv()
result = gzip.decompress(compressData).decode('utf-8')
if result[:7] == '{"ping"':
ts = result[8:21]
pong = '{"pong":' + ts + '}'
ws.send(pong)
ws.send(send_info)
# else:
# # print(result)
# with open('./test_result.txt', 'a') as f:
# #f.write(threading.currentThread().name + '\n')
# f.write(result + '\n')
except Exception as e:
print("error is:",e)
class ApiUser(WebSocketLocust):
task_set = UserBehavior
min_wait = 100
max_wait = 200
完整代码2：
# -*- encoding:utf-8 -*-
import gzip
import json
import random
import threading
import time
import zlib
from threading import Timer
import websocket
from gevent._semaphore import Semaphore
from locust import TaskSet, task, Locust, events
# TODO: 设置集合点...
all_locusts_spawned = Semaphore()
all_locusts_spawned.acquire()
def on_hatch_complete(**kwargs):
all_locusts_spawned.release()
events.hatch_complete += on_hatch_complete
t2 = 0
repCount = 0
sendCount = 0
pingCount = 0
stSend = 0
openTime = 0
reqLen = 0
recordSt = 0
repList = []
printCount = 1
reqSentCount = 1
symbols = ["etcusdt"]
subbedCount = 0
retSubTopicCount = 0
testFlag = 0
def on_message(ws, message):
global t2
global repCount
global sendCount
global pingCount
global stSend
global printCount
global reqList
global recordSt
global subbedCount
global retSubTopicCount
global reqSentCount
req_list = {
"req_str1": '{"req": "market.%s.kline.1min"}' % random.choice(symbols),
"req_str2": '{"req": "market.%s.depth.step0"}' % random.choice(symbols),
"req_str3": '{"req": "market.%s.trade.detail"}' % random.choice(symbols),
"req_str4": '{"req": "market.%s.detail"}' % random.choice(symbols),
# "req_str5": '{"req": "market.overview"}',
}
# 对返回来的压缩数据进行解压缩
ws_result = zlib.decompressobj(31).decompress(message)
result = json.loads(ws_result.decode('utf-8'))
print(result)
recordEd = time.time() # 为了判断什么时候统计数据的结束时间
recordCost = round((recordEd - recordSt) * 1000, 3) # 统计的结束时间减去统计的开始时间
# print(result)
if 'subbed' in result:
subbedCount = subbedCount + 1
if subbedCount % 5 == 0:
print("----------------subbed all topic----------------")
if 'ch' in result:
retSubTopicCount = retSubTopicCount + 1
if 'rep' in result:
repCount = repCount + 1
repRetTime = int((time.time() - stSend) * 1000)
repList.append(repRetTime)
# print("the server rep time is ---->%dms" % repRetTime)
# print("the server rep data is ---->%s" % result)
# 判断ping的返回 ，对应给服务器发送pong
if 'ping' in result:
pingCount = pingCount + 1
ping_id = result.get('ping')
pong_str = '{"pong": %d}' % ping_id
ws.send(pong_str)
t1 = ping_id
t3 = ping_id - t2
t2 = t1
if t3 > 5000:
print("$$$$$$$time difference ping is %d$$$$$$$ " % t3)
# print("ret ping value %d" % ping_id)
# print("ret ping curTime %d" % int(time.time()*1000))
# if 1000 < int((time.time()*1000) - ping_id):
# print("cur - pingTime is ---> %dms" % int((time.time()*1000) - ping_id))
if recordCost >= (random.randint(2000, 3000) * reqSentCount):
reqSentCount += 1
for key in req_list.keys():
ws.send(req_list[key])
sendCount = sendCount + 1
# print("send req info is --------->", req_list[key])
stSend = time.time()
# print("**********send count is %d *************** " % sendCount)
# 每1分钟统计一次
if recordCost >= (60000 * printCount):
printCount = printCount + 1
curTime = time.strftime('%Y-%m-%d %H:%M:%S')
repList.sort()
retCount = len(repList)
writeData = '| 当前时间%s ,req发送条数%s,返回总数据条数%s | 数据95耗时:%s | 数据50耗时:%s | sub返回量:%s ' % (
curTime, sendCount, repCount, repList[int(retCount * 0.95)], repList[int(retCount * 0.5)], retSubTopicCount)
fid = open("GipRecord.txt", "a+")
fid.write(writeData + "\n")
fid.close()
# 重新实现对应事件
def on_error(ws, error):
print("occur error " + error)
def on_close(ws):
global printCount
global reqSentCount
printCount = 1
reqSentCount = 1
print("^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^con is closed^^^^^^^^^^^^^^^^^^^^^^^^^^^^^")
def on_open(ws):
print("con success ...")
global reqList
global sendCount
global reqLen
global recordSt
global stSend
recordSt = time.time() # 为了统计记录文件创建的开始时间
stSend = time.time()
sub_list = {
"sub_str1": '{"sub": "market.%s.kline.1min"}' % random.choice(symbols),
"sub_str2": '{"sub": "market.%s.depth.step0"}' % random.choice(symbols),
"sub_str3": '{"sub": "market.%s.trade.detail"}' % random.choice(symbols),
"sub_str4": '{"sub": "market.%s.detail"}' % random.choice(symbols),
"sub_str5": '{"sub": "market.overview"}',
}
for key in sub_list.keys():
ws.send(sub_list[key])
class WSClient(object):
def __init__(self, host):
self.ws = None
self.host = host
def useWebCreate(self):
# websocket.enableTrace(True)
self.ws = websocket.WebSocketApp(self.host,
# header={'cloud-exchange':'510a02991'},
on_message=on_message,
on_error=on_error,
on_close=on_close,
on_open=on_open)
def execute(self):
self.ws.run_forever()
class AbstractLocust(Locust):
def __init__(self, *args, **kwargs):
super(AbstractLocust, self).__init__(*args, **kwargs)
self.client = WSClient(self.host)
class ApiUser(AbstractLocust):
host = 'ws://xxx/ws'
min_wait = 10
max_wait = 1000
class task_set(TaskSet):
def on_start(self):
self.client.useWebCreate()
# TODO: 设置集合点...
all_locusts_spawned.wait()
@task
def execute_long_run(self):
self.client.execute()
然后通过locust命令执行
locust -f xx.py --no-web -c 2 -r 1 -t 1m