Introduction

AWS’ Boto library is used commonly to integrate Python applications with various AWS services. I have generally avoided writing unit-tests for application code that interacts with the boto library because of the complexity involved in mocking and testing these functions.

However, I recently tried out the Moto library which makes it easy to mock AWS services and test code that interacts with AWS.

Some of the benefits of using Moto:

  • Testing code that interacts with AWS. Instead of having to test your code in an AWS environment, test AWS interactions locally.
  • Easy to learn and get started with.
  • Extensive coverage of AWS services.

In this article, we will look at how to add unit tests for AWS KMS using Moto.

Prerequisites

We need to install the following libraries to run the application:

  • boto3: AWS SDK for Python
  • moto: Mock AWS Services
  • Pytest: Test framework for Python

These libraries can be installed via pip:

pip install boto3 moto pytest

Getting Started

We will test the functionality that we added in this AWS KMS tutorial.

To test our KMS code, let’s create a file called test_kms.py with the following content:

import base64
import boto3
import pytest

from moto import mock_kms


@mock_kms
def test_list_keys():
    conn = boto3.client("kms","us-west-2")
    key = conn.create_key(Description="my key")

    response = conn.list_keys()
    assert len(response["Keys"]) == 1


@mock_kms
def test_describe_key():
    conn = boto3.client("kms","us-west-2")
    key = conn.create_key(Description="my key")
    key_id = key["KeyMetadata"]["KeyId"]

    key = conn.describe_key(KeyId=key_id)
    assert (
        key["KeyMetadata"]["Description"] == "my key"
    )


@mock_kms
def test_generate_data_key():
    conn = boto3.client("kms","us-west-2")
    key = conn.create_key(Description="my key")
    key_id = key["KeyMetadata"]["KeyId"]

    data_key = conn.generate_data_key(
        KeyId=key_id, NumberOfBytes=32,
    )
    # Plaintext must NOT be base64-encoded
    with pytest.raises(Exception):
        base64.b64decode(data_key["Plaintext"], validate=True)


@mock_kms
def test_decrypt_data_key():
    conn = boto3.client("kms","us-west-2")
    key = conn.create_key(Description="my key")
    key_id = key["KeyMetadata"]["KeyId"]

    data_key = conn.generate_data_key(
        KeyId=key_id, NumberOfBytes=32,
    )

    encrypted_data_key, plaintext_data_key = (
        data_key["CiphertextBlob"], data_key["Plaintext"]
    )
    response = conn.decrypt(CiphertextBlob=encrypted_data_key)
    assert response["Plaintext"] == plaintext_data_key

Things to note:

  • test_list_keys: In this test, we assert that the list of keys returns the key we just created.
  • test_describe_key: In this test, we assert that they key we created has the expected attributes.
  • test_generate_data_key: In this test, we check that we can generate a data key using the CMK we created earlier.
  • test_decrypt_data_key: In this test, we assert that the encrypted data key can be decrypted correctly.

Testing

We can run the unit tests by running the following command:


pytest test_kms.py
=============================================== test session starts ===============================================
platform linux -- Python 3.7.3, pytest-6.1.2, py-1.9.0, pluggy-0.13.1
rootdir: /home/abhishek/Projects/aws-projects/aws-boto-unit-tests
collected 4 items

test_kms.py ....                                                                                            [100%]

================================================ warnings summary =================================================
test_kms.py::test_list_keys
  /home/abhishek/.virtualenvs/aws-projects/lib/python3.7/site-packages/boto/plugin.py:40: DeprecationWarning: the imp module is deprecated in favour of importlib; see the module's documentation for alternative uses
    import imp

-- Docs: https://docs.pytest.org/en/stable/warnings.html
========================================== 4 passed, 1 warning in 0.51s ===========================================

Conclusion

Moto and Pytest make it pretty easy to unit test AWS-specific application code. It has already helped me become more confident about code that is being deployed to production and makes it possible to test such locally instead of in an AWS environment.