locust测试框架测试rabbirmq性能案例

locust测试框架默认为测试http请求，具体详细信息可参考其官方文档：
https://docs.locust.io/en/stable/what-is-locust.html
本片博客主要讲解locust测试自定义工具，以测试rabbitmq为例，话不多说上代码：
若读者还不清楚rabbitmq的工作机制，可参考我的另一篇博客，讲解rabbitmq的实现过程。

#!/usr/bin/env python
# coding:utf-8
import threading
import pika
import logging
import traceback
import uuid
import json
from pika.exceptions import ConnectionClosed
from time import sleep, time
from locust import Locust, TaskSet, task, events

class Logger:
    def __init__(self, path, clevel=logging.DEBUG, Flevel=logging.DEBUG):
        self.logger = logging.getLogger(path)
        self.logger.setLevel(logging.DEBUG)
        fmt = logging.Formatter('[%(asctime)s] [%(levelname)s] %(message)s', '%Y-%m-%d %H:%M:%S')
        # 设置CMD日志
        sh = logging.StreamHandler()
        sh.setFormatter(fmt)
        sh.setLevel(clevel)
        # 设置文件日志
        fh = logging.FileHandler(path)
        fh.setFormatter(fmt)
        fh.setLevel(Flevel)
        self.logger.addHandler(sh)
        self.logger.addHandler(fh)

    def debug(self, message):
        self.logger.debug(message)

    def info(self, message):
        self.logger.info(message)

    def war(self, message):
        self.logger.warn(message)

    def error(self, message):
        self.logger.error(message)

    def cri(self, message):
        self.logger.critical(message)

logger = Logger('hearbeat_return.log')


class RpcClient(object):
    """Asynchronous Rpc client."""
    _queue = {}
    _internal_lock = threading.Lock()

    def __init__(self, c_id, host="127.0.0.1", port=5672, username="admin", password="123456", v_host="/"):
        self._exchange_name = "istream_desktop"
        self._callback_queue_name = "rpc-client-{}".format(c_id)
        self._routing_key = "istream.rpc.client.{}".format(c_id)

        self._broker_url = "amqp://{}:{}@{}:{}{}".format(username, password, host, port, v_host)
        self._param = pika.URLParameters(self._broker_url)

        self._thread = threading.Thread(target=self.__run)
        self._thread.setDaemon(True)
        self._thread.start()
        self.start_time = None
        self.end_time = None

    def __run(self):
        while True:
            try:
                self.__connect()
                self.__process_data_events()
            except ConnectionClosed as e:
                logger.error("__run connection closed exception: {}".format(traceback.format_exc()))
                sleep(5)

    def __connect(self):
        self._connection = pika.BlockingConnection(parameters=self._param)
        self._channel = self._connection.channel()

        self._channel.exchange_declare(self._exchange_name, exchange_type='topic', durable=True, auto_delete=False)
        self._channel.queue_declare(self._callback_queue_name, exclusive=False, auto_delete=True)
        self._channel.queue_bind(self._callback_queue_name, self._exchange_name, routing_key=self._routing_key)

        self._thread = threading.Thread(target=self.__process_data_events)
        self._thread.setDaemon(True)
        self._thread.start()
        print(self._channel)

    def __process_data_events(self):
        self._channel.basic_consume(self.__on_response,queue=self._callback_queue_name,no_ack=True)
        while True:
            with self._internal_lock:
                self._connection.process_data_events()  # TODO: 连接关闭
            sleep(0.1)

    def __on_response(self, ch, method, props, body):
        self.end_time = time()
        print("start_time is {}".format(self.start_time))
        self._queue[props.correlation_id] = body
        body_json = json.loads(body)
        total_time = int((self.end_time - self.start_time) * 1000)
        if body_json.get('status') == 0:
            events.request_success.fire(request_type="send_hearbeat",name="wp",response_time=total_time,response_length=0)
        else:
            events.request_failure.fire(request_type="send_hearbeat",name="wp",response_time=total_time,exception=body_json.get('failed'))

    def __send_request(self, to, payload, notify=False):
        print("send is {}".format(self._channel))
        correlation_id = str(uuid.uuid4())

        props = pika.BasicProperties(content_type="application/json")

        if not notify:
            self._queue[correlation_id] = None
            props.reply_to = self._routing_key
            props.correlation_id = correlation_id

        with self._internal_lock:
            self.start_time = time()
            self._channel.basic_publish(exchange=self._exchange_name,
                                        routing_key="istream.{}".format(to),
                                        properties=props,
                                        body=str(payload))


        return correlation_id

    def call(self, to, func, params, timeout=10.0, notify=False):
        """调用远程接口，带返回内容

        :param to: 被调用方的 ID 标识, service.idata、 client.5f0c786d-2f38-4fd2-b081-ad7244a13e76、...
        :param func: 调用的接口名称
        :param params: 参数，字典
        :param timeout: 超时时间，默认10秒
        :param notify: 是否为通知消息(不等待响应)，默认为False
        :return: 接口处理响应内容
        """
        def current_milli_time():
            return int(time() * 1000)

        for p in (to, func):
            if not (isinstance(p, str)):
                logger.debug("status = 1001,message = 'RPC Parameter error'")

        if not (isinstance(params, dict) and isinstance(timeout, (int, float))):
            logger.debug("status = 1001,message = 'RPC Parameter error'")

        begin = current_milli_time()
        milli_timeout = timeout * 1000
        body = {"func": func, "params": params}
        correlation_id = self.__send_request(to, json.dumps(body), notify)

        if notify:
            """如果仅发送通知，则不等待接收响应消息"""
            return True

        while self._queue[correlation_id] is None:
            sleep(0.01)  # So, 精度 0.01 秒
            if current_milli_time() - begin > milli_timeout:
                logger.debug("status = 1001,message = 'RPC timeout'")
                break

        return self._queue.pop(correlation_id)

    def cast(self, to, func, params, timeout=10.0):
        return self.call(to, func, params, timeout, notify=True)



class Task_heartbeat(TaskSet):
    '''
    测试任务类
    '''
    c_id = "test_login"
    def __init__(self,parent):
        super(Task_heartbeat, self).__init__(parent)
        self.rpc_client = RpcClient(c_id=self.c_id,host='192.168.1.80')
        self.to_rpc = 'service.icache'

    @task(1)    #被task装饰的方法为被测试的方法，括号中的数字为被优先执行的权重，数字越小权重越大
    def send_heartbeat(self):
        fun_name = 'heartbeat'
        send_msg = {'terminal_id': "111111111", 'username': '222222', 'desktop': '333333'}
        body = self.rpc_client.call(self.to_rpc,fun_name,send_msg)
        logger.debug("heartbeat return body is {}".format(body))

    @task(2)
    def send_login(self):
        fun_name = 'login'
        send_msg = {}   #TODO:send_msg定义具体的请求参数
        body = self.rpc_client.call(self.to_rpc, fun_name, send_msg)
        logger.debug("login return body is {}".format(body))

    @task(3)
    def send_logout(self):
        fun_name = 'logout'
        send_msg = {}   #TODO:send_msg定义具体的请求参数
        body = self.rpc_client.call(self.to_rpc, fun_name, send_msg)
        logger.debug("login return body is {}".format(body))

class TestRabbitMq(Locust):
    '''
    locust测试类，测试入口，用于指定测试任务类
    '''
    min_wait = 100
    max_wait = 500
    task_set = Task_heartbeat

以上代码只有rabbitmq的client端，若需要实现其通讯需实现其server，具体实现过程可参考我的另一篇博客