Python3中使用locust 2.15.1压测grpc接口

[本文出自天外归云的博客园]

首先安装grpc、日志、locust相关依赖库：

google
protobuf
grpcio
grpcio-tools
grpc_interceptor
loguru
locust==2.15.1

然后创建一个grpc_user.py文件，内容如下：

import time
from typing import Any, Callable

import grpc
import grpc.experimental.gevent as grpc_gevent
from grpc_interceptor import ClientInterceptor
from locust import User
from locust.exception import LocustError
from loguru import logger

# patch grpc so that it uses gevent instead of asyncio
grpc_gevent.init_gevent()


class LocustInterceptor(ClientInterceptor):
    def __init__(self, environment, *args, **kwargs):
        super().__init__(*args, **kwargs)

        self.env = environment

    def intercept(
        self,
        method: Callable,
        request_or_iterator: Any,
        call_details: grpc.ClientCallDetails,
    ):
        response = None
        exception = None
        start_perf_counter = time.perf_counter()
        response_length = 0
        try:
            response = method(request_or_iterator, call_details)
            response_result = response.result()
            if response_result.code != 0:
                exception = Exception("code is not 0")
            else:
                # logger.debug(response_result)
                logger.debug("success")
            response_length = response.result().ByteSize()
        except grpc.RpcError as e:
            exception = e

        self.env.events.request.fire(
            request_type="grpc",
            name=call_details.method,
            response_time=(time.perf_counter() - start_perf_counter) * 1000,
            response_length=response_length,
            response=response,
            context=None,
            exception=exception,
        )
        return response


class GrpcUser(User):
    abstract = True
    stub_class = None

    def __init__(self, environment):
        super().__init__(environment)
        for attr_value, attr_name in (
            (self.host, "host"),
            (self.stub_class, "stub_class"),
        ):
            if attr_value is None:
                raise LocustError(f"You must specify the {attr_name}.")

        self._channel = grpc.insecure_channel(self.host)
        interceptor = LocustInterceptor(environment=environment)
        self._channel = grpc.intercept_channel(self._channel, interceptor)

        self.stub = self.stub_class(self._channel)

最后写压测文件（引用grpc_user.py），这里文件名假设为xxx.py。内容如下：

from locust import task

import xx_pb2, xx_pb2_grpc

import grpc_user


class MyGrpcUser(grpc_user.GrpcUser):
    host = "ip:port"
    stub_class = xx_pb2_grpc.XxxStub

    @task
    def get_charging_gun_count(self):
        req = xx_pb2.Req()
        self.stub.XxxMethod(req)