Locust - clients.py模块源码剖析

clients.py模块中核心包含了两个类，分别是HttpSession和ResponseContextManager，两者分别是对requests.Session和requests.Response进行封装得到的。

HttpSession类

class HttpSession(requests.Session):
    """
    HttpSession是基于requests.Session扩展的。
    主要用于执行Web请求并记录Session。
    每一个请求的相关信息都会进行记录，用于在最终的统计报告中进行展示。
    
    由于其本身是基于requests.Session扩展的，所以本身requests.Session的相关操作基本都可以正常支持。
    唯一的限制是请求的接口必须是与之前传入的HttpSession.base_url相匹配。
    
    HttpSession中进行请求时，需要额外传入两个参数（这是本身requests.Session中并不需要的）
    1. name：可以用于标记接口，而不是仅仅使用默认的URL path路径。它可以帮助我们把不同的url整个到同一个接口静态数据中。
    2. catch_response：一个布尔值。当设置为True时，接口的返回值会是一个context manager。通过context manager，我们可以更加明确的对成功、失败进行标记，而不仅仅是默认的返回值为2**是标记为正确，当返回值为4**或5**时标记为错误。
    """
    def __init__(self, base_url, *args, **kwargs):
        super(HttpSession, self).__init__(*args, **kwargs)
 
        self.base_url = base_url
        
        # 检查相关的基本auth信息，仅限于使用username和password登录的相关接口
        parsed_url = urlparse(self.base_url)
        if parsed_url.username and parsed_url.password:
            netloc = parsed_url.hostname
            if parsed_url.port:
                netloc += ":%d" % parsed_url.port
            
            # 从base_url中移除username和password
            self.base_url = urlunparse((parsed_url.scheme, netloc, parsed_url.path, parsed_url.params, parsed_url.query, parsed_url.fragment))
            # 使用basic auth进行请求
            self.auth = HTTPBasicAuth(parsed_url.username, parsed_url.password)
    
    def _build_url(self, path):
        # 校验请求url是否符合http请求格式，如果符合之间返回，不符合的话默认为传递的是相对路径，将base_url与请求url进行连接。
        if absolute_http_url_regexp.match(path):
            return path
        else:
            return "%s%s" % (self.base_url, path)
    
    def request(self, method, url, name=None, catch_response=False, **kwargs):
        """
        发送一个请求，其返回值为一个requests.Response对象。
        :param method: 请求类型
        :param url: 请求url
        :param name: (optional) 请求名称，用于生成统计报告
        :param catch_response: (optional) 是否记录响应，可以自主表单响应是否正确
        :param params: (optional) 字典或字符串
        :param data: (optional) 字典或字符串
        :param headers: (optional) 字典
        :param cookies: (optional) 字典或CookieJar对象
        :param files: (optional) 包含文件名等信息的字典（multipart encoding upload）.
        :param auth: (optional) Auth tuple or callable to enable Basic/Digest/Custom HTTP Auth.
        :param timeout: (optional) 浮点型数字
        :param allow_redirects: (optional) 布尔型，默认为True
        :param proxies: (optional) 字典
        :param stream: (optional) 布尔型，默认为False
        :param verify: (optional) 如果设置为True，将会进行ssl认证，默认为True
        :param cert: (optional) 字符串(.pem)或元组('cert', 'key')
        """
        # 获取真实需要请求的url
        url = self._build_url(url)
        
        # 需要存储在统计数据中（用于生成报告）的元数据
        request_meta = {}
        
        # 将请求之前的信息首先进行存储
        request_meta["method"] = method
        request_meta["start_time"] = time.time()
        
        response = self._send_request_safe_mode(method, url, **kwargs)
        
        # 记录消耗的时间
        request_meta["response_time"] = int((time.time() - request_meta["start_time"]) * 1000)
        
        # 记录请求的名称
        request_meta["name"] = name or (response.history and response.history[0] or response).request.path_url
        
        # 获取返回值中body的content的大小
        if kwargs.get("stream", False):
            request_meta["content_size"] = int(response.headers.get("content-length") or 0)
        else:
            request_meta["content_size"] = len(response.content or "")
        
        if catch_response:
            # 如果设置了catch_response，那么将会请求元信息进行保存并返回一个ContextManager
            response.locust_request_meta = request_meta
            return ResponseContextManager(response)
        else:
            # 否则直接判断是否引起异常，并触发相应的成功或失败事件。并返回requests.Response
            try:
                response.raise_for_status()
            except RequestException as e:
                events.request_failure.fire(
                    request_type=request_meta["method"], 
                    name=request_meta["name"], 
                    response_time=request_meta["response_time"], 
                    exception=e, 
                )
            else:
                events.request_success.fire(
                    request_type=request_meta["method"],
                    name=request_meta["name"],
                    response_time=request_meta["response_time"],
                    response_length=request_meta["content_size"],
                )
            return response
    
    def _send_request_safe_mode(self, method, url, **kwargs):
        """
        以安全模式发送请求，捕获其中可能是由于网络连接引起的错误。
        """
        try:
            return requests.Session.request(self, method, url, **kwargs)
        except (MissingSchema, InvalidSchema, InvalidURL):
            raise
        except RequestException as e: # 记录这些可能由于网络连接引起的错误
            r = LocustResponse()
            r.error = e
            r.status_code = 0  # 设置返回码为0，默认为null
            r.request = Request(method, url).prepare() 
            return r

ResponseContextManager类

class LocustResponse(Response):
    def raise_for_status(self):  
        # 判断是否有错误信息，如果有错误信息则引发异常
        if hasattr(self, 'error') and self.error:
            raise self.error
        Response.raise_for_status(self)
 
 
class ResponseContextManager(LocustResponse):
    """
    Response上下文管理器。
    主要可以提供重新标定请求是否成功或者失败，而不是按照默认的验证码进行判断。
    ResponseContextManager是在requests.Response上扩展的，新增了两个方法success和failure。
    """
    _is_reported = False
    
    def __init__(self, response):
        # 复制传入的响应数据
        self.__dict__ = response.__dict__
    
    def __enter__(self):
        return self
    
    def __exit__(self, exc, value, traceback):
        if self._is_reported:
            # 如果已经重新进行了成功或失败的标记，那么可以忽略默认的响应码
            return exc is None
        if exc:
            if isinstance(value, ResponseError):
                self.failure(value)
            else:
                return False
        else:
            try:
                self.raise_for_status()
            except requests.exceptions.RequestException as e:
                self.failure(e)
            else:
                self.success()
        return True
    
    def success(self):
        """
        用于标定响应结果是正确的。
        示例:
        with self.client.get("/does/not/exist", catch_response=True) as response:
            if response.status_code == 404:
                response.success()
        """
        events.request_success.fire(
            request_type=self.locust_request_meta["method"],
            name=self.locust_request_meta["name"],
            response_time=self.locust_request_meta["response_time"],
            response_length=self.locust_request_meta["content_size"],
        )
        self._is_reported = True
    
    def failure(self, exc):
        """
        用于标定响应结果是错误的。
        示例:
        with self.client.get("/", catch_response=True) as response:
            if response.content == "":
                response.failure("No data")
        """
        if isinstance(exc, six.string_types):
            exc = CatchResponseError(exc)
        
        events.request_failure.fire(
            request_type=self.locust_request_meta["method"],
            name=self.locust_request_meta["name"],
            response_time=self.locust_request_meta["response_time"],
            exception=exc,
        )
        self._is_reported = True