1. # -*- coding: utf-8 -*-#
    2. # file:         socketio
    3. # Author:       ShunZhe
    4. # Date:         2021/6/26
    7. import json
    8. import logging
    9. import re
    10. import time
    12. import gevent
    13. from websocket import create_connection
    14. from locust import User
    16. """自定义websocket客户端"""
    18. header = {
    19.     'Accept-Encoding': 'gzip, deflate, br',
    20.     'Connection': 'keep-alive',
    21.     'Content-Type': 'application/json;charset=UTF-8',
    22.     'User-Agent': 'Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) '
    23.                   'Chrome/87.0.4280.88 Safari/537.36 '
    25. }
    28. class SocketIO(object):
    29.     _locust_environment = None
    30.     start_time = 0
    32.     def __init__(self):
    33.         self.events = None
    34.         self.ws = None
    36.     def _require_failure(self, name, response_time, message, action, context, **kwargs):
    37.         """统计失败请求数"""
    38.         self.events.request_failure.fire(
    39.             request_type="接收数据",
    40.             name=name,
    41.             response_time=response_time,
    42.             response_length=len(message),
    43.             exception=f'{action} Content Error!',
    44.             context=context,
    45.         )
    47.     def _require_success(self, name, response_time, message, **kwargs):
    48.         """统计成功请求数"""
    49.         self.events.request_success.fire(
    50.             request_type="接收数据",
    51.             name=name,
    52.             response_time=response_time,
    53.             response_length=len(message),
    54.         )
    56.     def _require_fire(self, name, response_time, message, context, **kwargs):
    57.         """通用统计请求数"""
    58.         self.events.request.fire(
    59.             request_type="接收数据",
    60.             name=name,
    61.             response_time=response_time,
    62.             response_length=len(message),
    63.             exception=None,
    64.             context=context,
    65.         )
    67.     def connect(self, host):
    68.         """建立连接"""
    69.         self.ws = create_connection(host, header=header)
    70.         gevent.spawn(self.receive)
    72.     def receive(self, context=None):
    73.         """接收事件下发数据"""
    74.         message_regex = re.compile(r"(\d*)(.*)")
    75.         while True:
    76.             message = self.ws.recv()
    77.             total_time = int((time.time() - self.start_time) * 1000)  # 统计响应时间
    78.             logging.info(f"WSR: {message}")
    79.             m = message_regex.match(message)
    80.             if m is None:
    81.                 raise Exception(f"got no matches in {message}")
    82.             code = m.group(1)  # 正则获取code
    83.             json_string = m.group(2)  # 正则获取下发数据
    84.             if code == "0":
    85.                 name = "0 open"
    86.             elif code == "3":
    87.                 name = "3 heartbeat"
    88.             elif code == "40":
    89.                 name = "40 message ok"
    90.             elif code == "42":
    91.                 """
    92.                 根据具体的业务场景统计请求的成功或者失败情况
    93.                 此处校验op操作如果length返回小于1则为失败
    94.                 """
    95.                 obj = json.loads(json_string)
    96.                 ts_type, payload = obj
    97.                 name = f"{code} {ts_type}"
    98.                 if 'join_ack' == ts_type:
    99.                     if payload['code'] == 1 or payload['code'] == 20000002:
    100.                         self._require_success(name, total_time, message)
    101.                     else:
    102.                         self._require_failure(name, total_time, message, ts_type, context)
    104.                 if 'operation_ack' == ts_type:
    105.                     """定义op操作成功、失败断言"""
    106.                     if payload['code'] == 1 and payload['length'] >= 1:
    107.                         self._require_success(name, total_time, message)
    108.                     else:
    109.                         self._require_failure(name, total_time, message, ts_type, context)
    111.                 elif 'set_doc_ack' == ts_type:
    112.                     """定义set_doc操作成功、失败断言"""
    113.                     if payload['code'] == 1:
    114.                         self._require_success(name, total_time, message)
    115.                     else:
    116.                         self._require_failure(name, total_time, message, ts_type, context)
    118.             else:
    119.                 logging.info(f"Received unexpected message: {message}")
    120.                 continue
    121.             if name in ('42 operation_ack', '42 set_doc_ack', '42 join_ack'):
    122.                 pass
    123.             else:
    124.                 self._require_fire(name, total_time, message, context)
    126.     def send(self, body, context=None):
    127.         """发送数据"""
    128.         # if context is None:
    129.         #     context = {}
    130.         # if body == "2":
    131.         #     action = "2 heartbeat"
    132.         # else:
    133.         #     m = re.compile(r"(\d*)(.*)").match(body)
    134.         #     assert m is not None
    135.         #     code = m.group(1)
    136.         #     action = m.group(2).split(',')[0].split('\"')[1]
    137.         #     action = f"{code} {action}"
    138.         # self._require_fire(name=action, response_time=None, message=body, context=context)
    139.         # logging.info(f"WSS: {body}")
    140.         self.start_time = time.time()  # 开始发送数据的时间
    141.         self.ws.send(body)
    143.     def sleep_with_heartbeat(self, seconds):
    144.         """模拟心跳"""
    145.         while seconds >= 0:
    146.             gevent.sleep(min(25, seconds))
    147.             seconds -= 25
    148.             self.send("2")
    149.             # logging.info('模拟心跳～')
    152. class SocketIOUser(User):
    153.     abstract = True
    155.     def __init__(self, *args, **kwargs):
    156.         super(SocketIOUser, self).__init__(*args, **kwargs)
    157.         self.client = SocketIO()
    158.         self.client._locust_environment = self.environment
    159.         self.client.events = self.environment.events

    测试脚本：

    1. # -*- coding: utf-8 -*-#
    2. # file:         websock_op
    3. # Author:       ShunZhe
    4. # Date:         2021/6/26
    5. import json
    6. import os
    7. import queue
    8. import uuid
    9. from gevent import sleep
    10. from gevent._semaphore import Semaphore, BoundedSemaphore
    11. from locust import task, TaskSet, events, constant
    12. from locust.log import setup_logging
    13. from locust.runners import MasterRunner
    15. from socket_client import SocketIOUser
    17. setup_logging("INFO", None)
    19. # all_locusts_spawned = Semaphore()
    20. all_locusts_spawned = BoundedSemaphore()
    21. all_locusts_spawned.acquire()
    24. @events.spawning_complete.add_listener
    25. def on_hatch_complete(**kwargs):
    26.     all_locusts_spawned.release()  # 创建集合点
    29. # events.spawning_complete.add_listener(on_hatch_complete)  # 1.0之后的写法
    32. work_ids = queue.Queue()
    35. @events.init_command_line_parser.add_listener
    36. def add_user_parser(parser, **kwargs):
    37.     """自定义locust命令参数"""
    38.     parser.add_argument(
    39.         "-fi",
    40.         "--first_index",
    41.         help="获取截取列表首索引"
    42.              "[fc:]",
    43.         default=0,
    44.         type=str
    45.     )
    46.     parser.add_argument(
    47.         "-li",
    48.         "--last_index",
    49.         help="获取截取列表末索引"
    50.              "[:li]",
    51.         default=0,
    52.         type=str
    53.     )
    54.     args = parser.parse_args()
    55.     if args.first_index:
    56.         os.environ['first_index'] = args.first_index
    57.     if args.last_index:
    58.         os.environ['last_index'] = args.last_index
    61. @events.init.add_listener
    62. def on_locust_init(environment, **kwargs):
    63.     print("当前环境变量first_index为： %s" % os.environ.get('first_index'))
    64.     print("当前环境变量last_index为： %s" % os.environ.get('last_index'))
    65.     work_id_data()
    66.     print('初始化测试数据完毕')
    67.     if isinstance(environment.runner, MasterRunner):
    68.         print("master节点执行")
    69.     else:
    70.         print("worker节点执行")
    73. def work_id_data():
    74.     """队列读取作品ID"""
    75.     # work_ids = queue.Queue()
    76.     with open('work_id.txt', 'r', encoding='utf-8') as f:
    77.         for i in f.readlines()[int(os.environ.get('first_index')):int(os.environ.get('last_index'))]:
    78.             work_ids.put_nowait(int(i.strip('\n')))
    79.     return work_ids
    82. def version_id_data():
    83.     """生成version_ids字典"""
    84.     version_id = []
    85.     for i in range(1000000):
    86.         version_id.append(i)
    87.     version_ids = dict(zip(version_id, version_id))
    88.     return version_ids
    91. def read_doc():
    92.     """获取doc字符串"""
    93.     with open('set_doc.json', 'r', encoding="utf-8") as f:
    94.         data = json.load(f)
    95.         data = json.dumps(data)
    96.         return ''.join(('42', data))
    99. def create_op_data():
    100.     """op操作-创建用户同步"""
    101.     with open('user_sync.json', 'r', encoding="utf-8") as f:
    102.         data = json.load(f)
    103.         data = json.dumps(data)
    104.         return ''.join(('42', data))
    107. class ApiUser(TaskSet):
    108.     """压测op操作"""
    109.     index = 1
    111.     @task(15)
    112.     def test_op(self):
    113.         """op操作"""
    114.         all_locusts_spawned.wait()  # 限制在所有用户准备完成前处于等待状态
    115.         self.client.send('42["operation",{'
    116.                          f'"revision":{self.user.version_id.get(self.index)}'
    117.                          ',"op":[[["extra","user_state","275127340","current_stage_actor",{"r":true,"i":'
    118.                          f'"{uuid.uuid1()}"'
    119.                          '}]]]}]')
    120.         self.index += 1
    122.     @task(1)
    123.     def test_sleep_with_heartbeat(self):
    124.         """心跳"""
    125.         self.client.sleep_with_heartbeat(24)
    128. class JoinUser(TaskSet):
    129.     """验证最大连接数"""
    130.     index = 1
    132.     @task(1)
    133.     def test_op(self):
    134.         """op操作"""
    135.         sleep(10)
    136.         all_locusts_spawned.wait()  # 限制在所有用户准备完成前处于等待状态
    137.         self.client.send('42["operation",{'
    138.                          f'"revision":{self.user.version_id.get(self.index)}'
    139.                          ',"op":[[["extra","user_state","275127340","current_stage_actor",{"r":true,"i":'
    140.                          f'"{uuid.uuid1()}"'
    141.                          '}]]]}]')
    142.         self.index += 1
    144.     @task(1)
    145.     def test_sleep_with_heartbeat(self):
    146.         """心跳"""
    147.         self.client.sleep_with_heartbeat(10)
    150. class WebSocketUser(SocketIOUser):
    151.     host = "wss://press-socketcoll"
    152.     # wait_time = between(0.2, 0.5)   # 设置任务执行完毕等待0.2~0.5
    153.     wait_time = constant(1)  # 设置任务执行等待1s后开始下一个任务
    154.     # wait_time = constant_pacing(1)  # 设置任务启动总等待时间1s
    155.     version_id = version_id_data()  # op操作版本号
    156.     doc_data = read_doc()  # doc json
    157.     user_op = create_op_data()  # 创建用户同步
    158.     tasks = {ApiUser}
    159.     # tasks = {JoinUser}  # 测试场景：最大连接数
    161.     def on_start(self):
    162.         work_id = work_ids.get()  # 队列获取作品ID
    163.         print('当前作品为： %s' % work_id)
    164.         # work_id = 3515541  # 队列获取作品ID
    165.         doc_data = self.doc_data  # 作品doc
    166.         user_op = self.user_op  # 用户同步
    167.         token = 'CDM5LTk2MjYwOWVhMTA0YSJ9.2k2QUywuWjjHu5P2I'
    168.         url = f'wss://press-socketcoll/?work_id={work_id}&stag=1&token={token}&EIO=3&transport=websocket'
    169.         self.client.connect(url)
    170.         self.client.send(f'42["join",{work_id}]')
    171.         sleep(1)
    172.         self.client.send(doc_data)
    173.         sleep(1)
    174.         # 创建用户同步
    175.         self.client.send(user_op)
    176.         sleep(1)
    179. if __name__ == "__main__":
    180.     file_name = os.path.abspath(__file__)
    181.     os.system(f'locust -f {file_name} --master --first_index 0 --last_index 1000')
    183. # # setup Environment and Runner
    184. # env = Environment(user_classes=[WebSocketUser])
    185. # env.create_local_runner()
    186. # env.events.request_success.fire()
    188. # # start a WebUI instance
    189. # env.create_web_ui("127.0.0.1", 8089)
    191. # # start a greenlet that periodically outputs the current stats
    192. # gevent.spawn(stats_printer(env.stats))
    193. # #
    194. # # # start a greenlet that save current stats to history
    195. # gevent.spawn(stats_history, env.runner)
    197. # # start the test
    198. # env.runner.start(1000, spawn_rate=30)
    200. # # in 60 seconds stop the runner
    201. # gevent.spawn_later(30, lambda: env.runner.quit())
    203. # # wait for the greenlets
    204. # env.runner.greenlet.join()
    206. # # stop the web server for good measures
    207. # env.web_ui.stop()

    运行结果：

    image.png
    参考locust插件项目：https://github.com/SvenskaSpel/locust-plugins