Boto3 and Python unittest.mock
Tags: aws, boto, python
Published: 2024-04-24 10:30:03

I start this post by saying I'm not a professional software developer, I work mainly in IT Operations, although I write especially for IAC and small lambdas functions. 

When developing a Lambda function most of the time I need to interact with AWS Services via the famous `boto3` library; `boto3` is a powerful library developed and maintained by AWS which provides a communication framework to interact with native AWS Cloud Services. 

Every time I struggle to mock the library. I even considered using the `moto` stubber but, I'm not happy with what is provided. I aim to 

1. Create a mock which can stub the original call
2. I want the mock to expose all the methods to check which parameters have been used during the calls, how many times it has been called, etc... all the stuff included within the [`Mock`](https://docs.python.org/3/library/unittest.mock.html) class 

In this example, we will mock `boto3` while it creates a client for `RDS`.

Consider the following code

```python
import boto3


class MyRDSManager:
    def __init__(self) -> None:

        self._rds_client = boto3.client("rds")

    def delete_db_cluster_snapshot(self) -> None:
        self._rds_client.delete_db_cluster_snapshot(DBClusterSnapshotIdentifier="ABC")

    def get_snapshots(self) -> list[dict]:
        snapshots = []

        paginator = self._rds_client.get_paginator("describe_db_cluster_snapshots")
        pages = paginator.paginate(DBClusterIdentifier="MyCluster")

        for page in pages:
            page_snapshots = page.get("DBClusterSnapshots")
            for snapshot in page_snapshots:
                snapshots.append(snapshot)

        return snapshots
```

To test the above class I developed the following tests 

```python
import pytest
from unittest.mock import Mock, patch
from datetime import datetime

from main import MyRDSManager


@pytest.fixture(scope="function")
def prepare_mock():
    with patch("main.boto3.client") as mock_boto_client:
        # the first thing to do is to set the return_value attribute as itself
        # this will return the mock itself when the code runs `boto3.client("rds")`
        mock_boto_client.return_value = mock_boto_client

        def mock_paginate_describe_db_cluster_snapshots(*args, **kwargs):
            snapshot_type = kwargs.get(
                "SnapshotType"
            )  # get all the parameter passed to the call

            return [{
                "DBClusterSnapshots": [
                {
                    "DBClusterSnapshotIdentifier": "ABC",
                    "SnapshotCreationTime": "2024-01-01",
                    "SnapshotType": "manual",
                }
            ]}]

        mock_boto_client.get_paginator = Mock()
        mock_paginator = Mock(return_value=None)
        mock_paginator.paginate = Mock(return_value=None)
        mock_paginator.paginate.side_effect = mock_paginate_describe_db_cluster_snapshots
        
        mock_boto_client.get_paginator.return_value = mock_paginator

        mock_boto_client.delete_db_cluster_snapshot = Mock(return_value=None)

        my_rds = MyRDSManager()

        yield my_rds, mock_boto_client


def test_one(prepare_mock):
    mock_my_rds, mock_boto_client = prepare_mock

    mock_my_rds.delete_db_cluster_snapshot()

    mock_boto_client.delete_db_cluster_snapshot.assert_called_once_with(DBClusterSnapshotIdentifier="ABC")

    result = mock_my_rds.get_snapshots()

    assert result == [{
        "DBClusterSnapshotIdentifier": "ABC",
        "SnapshotCreationTime": "2024-01-01",
        "SnapshotType": "manual",
    }]
```

The first thing to notice is that we need to set the ` mock_boto_client.return_value` to `mock_boto_client`, which is itself, and this will return the mock **instance** you are configuring when in `MyRDSManager` the code runs `self._rds_client = boto3.client("rds")`. If you don't set this MagicMock will return the default value, therefore a **new** MagicMock, and not what you are configuring!

Then you want to configure the Mock as you please, adding methods and attributes. Is worth noticing I set a `side_effect`: this will allow the mock to call the function I set in the `side_effect`, passing to it all the arguments the code is passing to the mock. 

With this setup, you should be able to fulfil the initial requirements and therefore stub the original behaviour and use all the Mock-provided features.

Hope this will help you all! 

Any feedback is appreciated. 

Cheers 

Post published also on Virtuability's website [here](https://www.virtuability.com/blog/2024-04-24-boto3-and-python-unittest-mock/).