Adding support for boto3 to Betamax testing framework

One of the API clients in Python that I really like is github3.py written by Ian Cordasco. It's one of my go-to places when I think about how to structure a Python API that I am building around a Web Service.

This is how I came across one of his other projects: betamax. It's a package that is specifically aimed at testing against an actual API such as github. He goes into more details in a recent talk he gave at PyCon US 2015 on the topic. It is a great piece of software that allows you to test against the API once, automatically record the request and response, and then replay it every time the tests are run again.

I won't go into any further detail on why this is a good idea because the talk mentioned above is a much better resource for that. What I want to write briefly about is that how betamax can be used in conjunction with boto3, the preview of Amazon's re-written API client.

boto3 internally uses requests to make call to the AWS API but make some very substantial customizations to it. The package also uses a considerable amount of dynamic class generation, enough to make understanding it more difficult than I would like it to be.

So long story short, I was trying to integrate betamax into my tests for a project that uses boto3. It took "a little while" to work out what happens where and how but I eventually found a very simple way to hook betamax into boto3.

The first step is to configure betamax. I am using pytest as my test runner so I'm adding it into my conftest.py:

# conftest.py
os.makedirs('tests/cassettes')

with betamax.Betamax.configure() as config:
    config.cassette_library_dir = 'tests/cassettes/'

    # Remove credentials from serialized requests and responses.
    config.define_cassette_placeholder(
        '<AWS_ACCESS_KEY_ID>', AWS_ACCESS_KEY_ID)
    config.define_cassette_placeholder(
        '<AWS_SECRET_ACCESS_KEY>', AWS_SECRET_ACCESS_KEY)

This setup is exactly as described in the great docs for betamax itself and is used in github3.py as well. The next thing we have to do is write a little wrapper around the recorder class betamax.Betamax that we use to take care of the recording of requests and responses. This is where the magic happens, and where you have to know how a resource generated with boto3 wraps a requests session. A resource in boto3 represents an API endpoint for a specific service such as EC2 and can be created using:

session = boto3.session.Session()
ec2_resource = session.resource('ec2')

The ec2_resource can now be used to interact with the API. To make it easy to use a betamax recorder with such a resource, we can create the following wrapper class:

# botomax.py
import betamax

class Botomax(betamax.Betamax):
    def __init__(self, boto_resource, cassette_library_dir=None,
                 default_cassette_options={}):

        # This is all the magic
        resource_session = boto_resource.meta.client._endpoint.http_session

        super(Botomax, self).__init__(
            session=resource_session,
            cassette_library_dir=cassette_library_dir,
            default_cassette_options=default_cassette_options)

And with this simple trick, I am now able the Botomax recorder in my tests integrating with the AWS API. An example test looks like this and will create a new file called example.json in the tests/cassettes/ directory. Each file is created once by making an actual API call. After that, the stored "cassettes" are re-played instead of making actual API calls.

# test_ec2.py
def test_ec2_connection_with_betamax():
    session = boto3.session.Session()
    ec2 = session.resource('ec2')

    with Botomax(ec2) as vcr:
        vcr.use_cassette('example')

        for instance in ec2.instances.all():
            print(instance)

I hope this was useful for you. If you have feedback or questions, you can tweet at me or get in touch through github.

Loading