DEV Community

Danilo Desole
Danilo Desole

Posted on • Edited on

Boto3 and Python unittest.mock

#aws #boto #python

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 class

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

Consider the following code

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
Enter fullscreen mode Exit fullscreen mode

To test the above class I developed the following tests

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",
    }]
Enter fullscreen mode Exit fullscreen mode

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!

Another important point to note is that we patch boto3 within the specific module being tested, rather than the general boto3 library. In other words, you should patch main.boto3.client, where main refers to the module you have written and are currently testing.

Next, configure the Mock as needed by adding methods and attributes. It's worth noting that setting a side_effect allows the mock to invoke the function specified in the side_effect, passing along all the arguments that the code supplies 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.

Top comments (3)

Collapse
 
mauricebrg profile image
Maurice Borgmeier •

Nice post :)

I'm curious how moto didn't work for you - are the API calls not supported (yet)?

Collapse
 
panilo profile image
Danilo Desole •

Hi Maurice! Thank you :) moto didn't work for me for 2 reasons

  1. I didn't want to configure the moto virtual AWS environment
  2. It doesn't give you the possibility to leverage Mock-related features like boto_mocked_function.assert_called_once_with(...) and so on...

Said that though I agree moto might fit most of the developers :)

Collapse
 
mauricebrg profile image
Maurice Borgmeier •

I see - yeah, moto is more suitable if you want to test the behavior on a slightly higher level, i.e. the assertion for delete_db_cluster_snapshot would be checking if the given snapshot has been removed from the list_db_cluster_snapshot. The other benefit of moto is that you get all the parameter validation the boto3 does when you use an API call.

I frequently use moto for integration tests where I want to verify the interaction of multiple components.