Locust and going deeper with Python Requests

I'm doing some work with Locust as a performance and load test tool which has driven me to learn Python. Like any knew language/technology, I set out to learn enough to solve my problem then circled back and trying to find better ways of doing what I did. When you're new to a technology, you're especially ignorant about what already exists that can be leveraged. Pro tip: this phase of your process is a good time to go back and read the docs (even if you already have). Now that you know a little bit about the technology, the docs will be a lot more meaningful to you.

Locust is based on the defacto standard http client in Python called Requests. After writing few Locust tests I realized I had built a lot of helper methods and general "wrapping" code around Requests in order to accomplish orthogonal things necessary for the API I was testing. For mature technologies (such as Python) there is almost always a better way to do such things. Below are some things I learned.

The API was testing needed a few basic orthogonal things:

header values returned in the response of one call needed to be captured in order to send as a header in the next request
post-processing after each request
logging mechanism in order to troubleshooting API calls

In reverse order...

Logging

Logging can easily be solved with the following code (borrowed from this SO post):

import requests
from requests.auth import AuthBase
import logging
import contextlib
try:
    from http.client import HTTPConnection # py3
except ImportError:
    from httplib import HTTPConnection # py2

def debug_requests_on():
    '''Switches on logging of the requests module.'''
    HTTPConnection.debuglevel = 1

    logging.basicConfig()
    logging.getLogger().setLevel(logging.DEBUG)
    requests_log = logging.getLogger("requests.packages.urllib3")
    requests_log.setLevel(logging.DEBUG)
    requests_log.propagate = True

debug_requests_on()

requests.get('https://google.com')

which produces:

$ requests-playing.py
DEBUG:urllib3.connectionpool:Starting new HTTPS connection (1): google.com:443
send: b'GET / HTTP/1.1\r\nHost: google.com\r\nUser-Agent: python-requests/2.22.0\r\nAccept-Encoding: gzip, deflate\r\nAccept: */*\r\nConnection: keep-alive\r\n\r\n'
reply: 'HTTP/1.1 301 Moved Permanently\r\n'
header: Location: https://www.google.com/
header: Content-Type: text/html; charset=UTF-8
header: Date: Sun, 01 Sep 2019 16:04:22 GMT
header: Expires: Tue, 01 Oct 2019 16:04:22 GMT
header: Cache-Control: public, max-age=2592000
header: Server: gws
header: Content-Length: 220
header: X-XSS-Protection: 0
header: X-Frame-Options: SAMEORIGIN
header: Alt-Svc: quic=":443"; ma=2592000; v="46,43,39"
DEBUG:urllib3.connectionpool:https://google.com:443 "GET / HTTP/1.1" 301 220
DEBUG:urllib3.connectionpool:Starting new HTTPS connection (1): www.google.com:443
send: b'GET / HTTP/1.1\r\nHost: www.google.com\r\nUser-Agent: python-requests/2.22.0\r\nAccept-Encoding: gzip, deflate\r\nAccept: */*\r\nConnection: keep-alive\r\n\r\n'
reply: 'HTTP/1.1 200 OK\r\n'
header: Date: Sun, 01 Sep 2019 16:04:22 GMT
header: Expires: -1
header: Cache-Control: private, max-age=0
header: Content-Type: text/html; charset=ISO-8859-1
header: P3P: CP="This is not a P3P policy! See g.co/p3phelp for more info."
header: Content-Encoding: gzip
header: Server: gws
header: X-XSS-Protection: 0
header: X-Frame-Options: SAMEORIGIN
header: Set-Cookie: 1P_JAR=2019-09-01-16; expires=Tue, 01-Oct-2019 16:04:22 GMT; path=/; domain=.google.com; SameSite=none
header: Set-Cookie: NID=188=KKTQZzXKgXv-VsCoEc6FSr0-Bw5xG3iA7X72tCvEq7cc7riQmniX4bZ-YlUH9BrByAv10-Shc62KHZXaL3QkQImJA6PmmqZog7i8vz-DMj-E7DUMNpqSSXYXnlmg4qW0SQhrkJ6A6X6Hti593PKd06ABjlUH-vtEWnfkLNROR3E; expires=Mon, 02-Mar-2020 16:04:22 GMT; path=/; domain=.google.com; HttpOnly
header: Alt-Svc: quic=":443"; ma=2592000; v="46,43,39"
header: Transfer-Encoding: chunked
DEBUG:urllib3.connectionpool:https://www.google.com:443 "GET / HTTP/1.1" 200 None

Post-Processing:

It turns out Requests supports a hooks mechanism. As best I can tell from the docs, they only support a single hook (`response`) so maybe it shouldn't be plural, but regardless, you can do this:

def call_hook(resp, *args, **kwargs):
    print(f"{resp.url} returned status code of {resp.status_code}")


requests.get('https://google.com', hooks={'response': call_hook})

which produces:

$ requests-playing.py
https://google.com/ returned status code of 301
https://www.google.com/ returned status code of 200

notice it's following the redirect automatically and the hook is respected for each. If that's undesirable, that can be suppressed with the allow_redirects flag set to False:

def call_hook(resp, *args, **kwargs):
    print(f"{resp.url} returned status code of {resp.status_code}")

requests.get('https://google.com', hooks={'response': call_hook}, allow_redirects=False )

which produces:

$ requests-playing.py
https://google.com/ returned status code of 301

Header Chaining

In conjunction with what's described above.... the Custom Authentication mechanism built into Requests gives me what I need

class MyAPIAuth(AuthBase):

    authorization_header = None

    def __init__(self, authorization_header):
        self.authorization_header = authorization_header

    def __call__(self, req):
        req.headers['authorization'] = self.authorization_header
        return req

requests.get('https://api.mysite.com', auth=MyAPIAuth('Basic xxxxxxxxxxxxx')

WARNING! That __call__ method must return the request object or you'll have weird problems :)