Python AWS testing with boto stubber
Tags: aws, python, testing, programming
Published: 2023-06-21 22:47:11

Making sure code does what it's supposed to is a crucial part of the development process. This is even more crucial when dealing with AWS code that can have unexpected costs if something goes wrong. Here I'll show how [boto Stubber](https://botocore.amazonaws.com/v1/documentation/api/latest/reference/stubber.html) can help with testing code.

## The Initial Setup

The boto library used for making python boto calls actually has built in functionality to help with testing. Due to wrapping around the client it's recommended to have any clients instantiated via a specific function/method to make it more easily mockable. Take for example a simple cleanup script that terminates EC2 instances:

```python
import boto3


class InstanceCleaner(object):
    def __init__(self) -> None:
        self.client = self.create_client()

    @staticmethod
    def create_client() -> boto3.client:
        return boto3.client('ec2', region_name='us-west-2')

    def terminate_instances(self) -> None:
        instances = self.client.describe_instances()
        instance_list = []
        if not instances['Reservations']:
            return False

        for instance in instances['Reservations'][0]['Instances']:
            instance_list.append(instance['InstanceId'])

        self.client.terminate_instances(InstanceIds=instance_list)
```

So nothing complicated, it's just collecting an list of instance IDs and passing them on to `terminate_instances`. Now I'll make a test suite  for this:

```python
from instance_cleanup import InstanceCleaner
from botocore.stub import Stubber

mock_describe_instances = {
    'Reservations': [
        {
            'Instances': [
                {
                    'InstanceId': 'i-1234567890abcdef0',
                    'State': {
                        'Code': 16,
                        'Name': 'running'
                    }
                },
            ],
        }
    ]
}

mock_terminate_instances = {
    'TerminatingInstances': [
        {
            'CurrentState': {
                'Code': 32,
                'Name': 'shutting-down',
            },
            'InstanceId': 'i-1234567890abcdef0',
            'PreviousState': {
                'Code': 16,
                'Name': 'running',
            },
        },
    ]
}

def test_instance_cleanup():
    instance_cleanup = InstanceCleaner()
    with Stubber(instance_cleanup.client) as stubber:
        stubber.add_response('describe_instances', mock_describe_instances, {})
        stubber.add_response('terminate_instances', mock_terminate_instances, {'InstanceIds': ['i-1234567890abcdef0']})
        instance_cleanup.terminate_instances()
```

So here there are two JSON mocks to handle the responses for both `describe_instances` and `terminate_instances`. In the case of `mock_describe_instances` it's not using the full JSON return available for the call. While you can use this to slim the mocks down, there are cases where boto can complain about missing specific fields in the return structure. It is possible to ignore such issues, but I generally like to stay with what the API expects. That way if the test fail due to a new required field in the results, I can investigate to see if other changes may effect my tests. Note that I'm using a [context manager](https://docs.python.org/3/library/contextlib.html) here for the stubbed client:

```python
with Stubber(instance_cleanup.client) as stubber:
```

The reason why is that I'd have to use `stubber.activate()` to register the responses which can be easy to miss. Using the context manager instead let's the system handle it transparently for me. A quick run of the tests gives the following:

```
> python -m pytest
=========================================================================================================================================================================== test session starts ===========================================================================================================================================================================
platform win32 -- Python 3.9.13, pytest-7.3.2, pluggy-1.0.0
rootdir: [redacted]
collected 1 item

tests\test_instance_cleanup.py .                                                                                                                                                                                                                                                                                                                                     [100%]

============================================================================================================================================================================ 1 passed in 0.48s ============================================================================================================================================================================ 
```

## Stubber Gotchas

This allows for testing AWS boto calls without actually connecting to the AWS API itself. It's also important to note that the order of `add_response` is crucial. For example if you changed the test to do this:

```python
def test_instance_cleanup():
    instance_cleanup = InstanceCleaner()
    with Stubber(instance_cleanup.client) as stubber:
        stubber.add_response('terminate_instances', mock_terminate_instances, {'InstanceIds': ['i-1234567890abcdef0']})
        stubber.add_response('describe_instances', mock_describe_instances, {})
        instance_cleanup.terminate_instances()
```

Then you'll get an error when trying to run the test:

```
>           raise StubResponseError(
                operation_name=model.name,
                reason=f'Operation mismatch: found response for {name}.',
            )
E           botocore.exceptions.StubResponseError: Error getting response stub for operation DescribeInstances: Operation mismatch: found response for TerminateInstances
```

Parameters also need to match what the call is made with. For example:

```python
def test_instance_cleanup():
    instance_cleanup = InstanceCleaner()
    with Stubber(instance_cleanup.client) as stubber:
        stubber.add_response('describe_instances', mock_describe_instances, {})
        stubber.add_response('terminate_instances', mock_terminate_instances, {})
        instance_cleanup.terminate_instances()
```

This will error out:

```
E           botocore.exceptions.StubAssertionError: Error getting response stub for operation TerminateInstances: Expected parameters:
E           {},
E           but received:
E           {'InstanceIds': ['i-1234567890abcdef0']}

..\.venv\lib\site-packages\botocore\stub.py:392: StubAssertionError
```

## State Mocking

As the core code is calling with a specific instance ID array so it won't match the blank parameters entry given. Now later we realize it would be more efficient to only attempt termination of instances that actually do need to be terminated, or those not in the terminated/about to be terminatd state. This leads to some code refactoring:

```python
import boto3

TERMINATABLE_STATES = ['pending', 'running', 'stopping', 'stopped']

class InstanceCleaner(object):
    def __init__(self) -> None:
        self.client = self.create_client()

    @staticmethod
    def create_client() -> boto3.client:
        return boto3.client('ec2', region_name='us-west-2')

    def terminate_instances(self) -> None:
        instances = self.client.describe_instances(
            Filters=[{
                'Name': 'instance-state-name',
                'Values': TERMINATABLE_STATES,
            }]
        )
        instance_list = []
        if not instances['Reservations']:
            return False

        for instance in instances['Reservations'][0]['Instances']:
            instance_list.append(instance['InstanceId'])

        if instance_list:
            self.client.terminate_instances(InstanceIds=instance_list)
```

Now if nothing is returned by `describe_instances` due to the filter then there is no need to make a termination call. Now for the test:

```python
import pytest

from instance_cleanup import InstanceCleaner, TERMINATABLE_STATES
from botocore.stub import Stubber

INSTANCE_CODE_MAPPING = {
    'pending': 0,
    'running': 16,
    'shutting-down': 32,
    'terminated': 48,
    'stopping': 64,
    'stopped': 80
}
INSTANCE_ID = 'i-1234567890abcdef0'

test_data = [(x, True) for x in INSTANCE_CODE_MAPPING.keys() if x != 'terminated' and x != 'shutting-down']
test_data.append(('shutting-down', False))
test_data.append(('terminated', False))

def generate_instance_response(code_name: str) -> dict:
    return {
        'Reservations': [
            {
                'Instances': [
                    {
                        'InstanceId': INSTANCE_ID,
                        'State': {
                            'Code': INSTANCE_CODE_MAPPING[code_name],
                            'Name': code_name
                        }
                    },
                ],
            }
        ]
    }

def generate_termination_response(code_name: str) -> dict:
    return {
        'TerminatingInstances': [
            {
                'CurrentState': {
                    'Code': 32,
                    'Name': 'shutting-down'
                },
                'InstanceId': INSTANCE_ID,
                'PreviousState': {
                    'Code': INSTANCE_CODE_MAPPING[code_name],
                    'Name': code_name,
                },
            },
        ]
    }

@pytest.mark.parametrize('code_name,should_terminate', test_data)
def test_instance_cleanup(code_name, should_terminate) -> None:
    instance_cleanup = InstanceCleaner()
    with Stubber(instance_cleanup.client) as stubber:
        if should_terminate:
            stubber.add_response('describe_instances', generate_instance_response(code_name), {'Filters':[{'Name': 'instance-state-name', 'Values': TERMINATABLE_STATES }]})
            stubber.add_response('terminate_instances', generate_termination_response(code_name), {'InstanceIds': [INSTANCE_ID]})
        else:
            stubber.add_response('describe_instances', {'Reservations': [{'Instances': []}]}, {'Filters':[{'Name': 'instance-state-name', 'Values': TERMINATABLE_STATES }]})
        instance_cleanup.terminate_instances()
```
So this is going through all the states that should be terminated and testing out how each one works. There's also one for shutting down (termination in progress) and terminated state (termination completed). Parameterized arguments are used to keep the test to a single function while being able to test multiple variations of state:

```python
@pytest.mark.parametrize('code_name,should_terminate', test_data)
def test_instance_cleanup(code_name, should_terminate) -> None:
```
in this case the first argument is the mapping to parameter names of `code_name` and `should_terminate`. The values for these are the pair of code name and value for each tuple in the `test_data` list:

```python
test_data = [(x, True) for x in INSTANCE_CODE_MAPPING.keys() if x != 'terminated' and x != 'shutting-down']
test_data.append(('shutting-down', False))
test_data.append(('terminated', False))
```

This data is generated off the keys of `INSTANCE_CODE_MAPPING` using a conditional [list comprehension](https://docs.python.org/3/tutorial/datastructures.html#list-comprehensions). Two functions are also declared to generate the appropriate mock JSON given the state in question. This could also be refactored out to a [pytest mock](https://docs.pytest.org/en/7.1.x/how-to/fixtures.html) as well.

## Conclusion

Boto stubber is an excellent way to handle testing. However it might not be the best use for cases like testing DynamoDB inserts and the JSON mocks can be difficult to manage in more complex cases. In the next installment I'll be looking at using moto as a testing alternative.