DEV Community

Syed Mohammad Ibrahim
Syed Mohammad Ibrahim

Posted on

Mocking AWS Dynamo Db calls during testing

Introduction

It is always a good idea to write test cases for your application when you build any new feature or the application from scratch itself. Testing allows your code to be more bug free and gives a clear sense of where the development should go and if there are any bottlenecks involved.

In this space, I am going to go over the basics of how to mock AWS Dynamo DB calls in while testing in Python. I will be using the moto library to do that.

Setup

I will be using Python version 3.9.x for the demonstration and following additional libraries:

  • pytest - version 7.1.1
  • moto - version 3.1.5
  • boto3 - version 1.21.45

I will be using virtualenv to manage my packages, but you are open to use whatever seems appropriate.

To install these libraries, use the following commands on your terminal or command prompt:

pip install pytest==7.1.1.
Enter fullscreen mode Exit fullscreen mode

Installs pytest version 7.1.1

pip install moto[all]==3.1.5
Enter fullscreen mode Exit fullscreen mode

Installs moto version 3.1.5 with all of it's dependencies. The moto library can mock a lot of AWS services. For installing all of it's dependencies - which it doesn't install by default - we need to add the [all] part while installing it.

pip install boto3==1.21.45
Enter fullscreen mode Exit fullscreen mode

Installs boto3 version 1.21.45

You can check out the complete list of requirements for this project on my GitHub page.

Project Structure

The project structure for this demo is going to look like this

MockDynamoDbExample
├── __init__.py
├── src
│   ├── __init__.py
│   └── service.py
└── test
    ├── __init__.py
    └── test_service.py
Enter fullscreen mode Exit fullscreen mode

Implementation

Let's start implementing the mocks.

First, create a base folder (in my case, it is MockDynamoDbExample). Then create two sub-folders src and test. As the name suggest, src will contain the code which needs to be tested and test will contain our tests to the given source code.

Next, we need a running service that will more or less reflect the operations we want to perform in the real world. Create empty __init__.py and service.py files in the src folder. The __init__.py is created so that the current folder is treated as a module. We are not going to edit that file. By just existing, the file indicates to the Python Interpreter that the current folder is a module. The service.py file is where we are going to write the source code that requires testing.

Let's start by writing all the imports and the constructor method

from boto3.dynamodb.conditions import Key
import boto3


class Service:
    __dynamodb = None
    __table = None

    def __init__(self, table_name: str, env="dev") -> None:
        """
        Initiates an object of this class

        :param table_name: The table this service is going to perform CRUD operations for.
        :param env: The environment context the application is running in.
        """

        # Create a boto3 resource for local instance of dynamodb
        resource = {"region_name": "us-east-2"}
        if env == "dev":
            resource["endpoint_url"] = "http://localhost:8000"
        self.__dynamodb = boto3.resource("dynamodb", **resource)

        # Get the table
        self.__get_table(table_name)
Enter fullscreen mode Exit fullscreen mode

The gist of the code in the constructor method is that it creates a boto3.resource instance depending on the env value passed. During testing, we won't be deploying an instance of Dynamo Db, that's why I had put an env check. You are free to play around with it.

Let's implement the __get_table method which is responsible for either getting us the table, if it exists or create a new one. Copy the below code in the same class

def __get_table(self, table_name: str):
    # Get the table from the dynamodb instance
    self.__table = self.__dynamodb.Table(table_name)

    try:
        # Check whether the table exists or not
        self.__table.table_status
    except:
        # If the table doesn't exist, an exception will be thrown.
        # Create the table if that happens.
        print("The table user doesn't exist. Creating now...")
        self.__create_table(table_name)

def __create_table(self, table_name: str):
    try:
        table = self.__dynamodb.create_table(
            TableName=table_name,
            KeySchema=[
                {"AttributeName": "username", "KeyType": "HASH"},
                {"AttributeName": "age", "KeyType": "RANGE"},
            ],
            AttributeDefinitions=[
                {
                    "AttributeName": "username",
                    "AttributeType": "S",
                },
                {
                    "AttributeName": "age",
                    "AttributeType": "N",
                },
            ],
            ProvisionedThroughput={
                "ReadCapacityUnits": 1,
                "WriteCapacityUnits": 1,
            },
        )

        # Wait until the table exits
        table.wait_until_exists()
    except Exception:
        raise Exception(f"Unable to create the table {table_name}")

    # Set the instance variable to the new table created
    self.__table = table

def delete_table(self):
    self.__table.delete()
Enter fullscreen mode Exit fullscreen mode

To give an overview, our table contains a primary-key username which is of key type HASH and age as a key type of RANGE. I am using the create_table method that is provided as part of the boto3.resource instance to create the desired table. The upstream code, namely __get_table method is responsible for checking whether the table exists in the first place or not. If it does, we return early with it's instance.

Now, we implement all the CRUD operations that are required to perform the task for the database.

def create_user(self, username: str, email: str, age: int, first_name: str, last_name: str):
    user_obj = {
        "username": username,
        "email": email,
        "age": age,
        "first_name": first_name,
        "last_name": last_name
    }

    # Insert the data into table
    self.__table.put_item(Item=user_obj)

def get_user_by_username(self, username: str):
    response = self.__table.query(
        KeyConditionExpression=Key("username").eq(username)
    )
    items = response["Items"]

    # Return none, if the length of the items list is not 1
    if len(items) != 1:
        return None

    # Return the value fetched from the database
    return items[0]

def delete_user_by_username(self, username: str):
    response = self.__table.delete_item(
        Key={"username": username}
    )

    return response
Enter fullscreen mode Exit fullscreen mode

The above CRUD operations can vary depending on your use case. For the purpose of this blog, I will be using these basic operations. You can implement your own CRUD operations by following this link.

The complete service.py file should look like

from boto3.dynamodb.conditions import Key
import boto3


class Service:
    __dynamodb = None
    __table = None

    def __init__(self, table_name: str, env="dev") -> None:
        """
        Initiates an object of this class

        :param table_name: The table this service is going to perform CRUD operations for.
        :param env: The environment context the application is running in.
        """

        # Create a boto3 resource for local instance of dynamodb
        resource = {"region_name": "us-east-2"}
        if env == "dev":
            resource["endpoint_url"] = "http://localhost:8000"
        self.__dynamodb = boto3.resource("dynamodb", **resource)

        # Get the table
        self.__get_table(table_name)

    def __get_table(self, table_name: str):
        # Get the table from the dynamodb instance
        self.__table = self.__dynamodb.Table(table_name)

        try:
            # Check whether the table exists or not
            self.__table.table_status
        except:
            # If the table doesn't exist, an exception will be thrown.
            # Create the table if that happens.
            print("The table user doesn't exist. Creating now...")
            self.__create_table(table_name)

    def __create_table(self, table_name: str):
        try:
            table = self.__dynamodb.create_table(
                TableName=table_name,
                KeySchema=[
                    {"AttributeName": "username", "KeyType": "HASH"},
                    {"AttributeName": "age", "KeyType": "RANGE"},
                ],
                AttributeDefinitions=[
                    {
                        "AttributeName": "username",
                        "AttributeType": "S",
                    },
                    {
                        "AttributeName": "age",
                        "AttributeType": "N",
                    },
                ],
                ProvisionedThroughput={
                    "ReadCapacityUnits": 1,
                    "WriteCapacityUnits": 1,
                },
            )

            # Wait until the table exits
            table.wait_until_exists()
        except Exception:
            raise Exception(f"Unable to create the table {table_name}")

        # Set the instance variable to the new table created
        self.__table = table

    def delete_table(self):
        self.__table.delete()

    def create_user(self, username: str, email: str, age: int, first_name: str, last_name: str):
        user_obj = {
            "username": username,
            "email": email,
            "age": age,
            "first_name": first_name,
            "last_name": last_name
        }

        # Insert the data into table
        self.__table.put_item(Item=user_obj)

    def get_user_by_username(self, username: str):
        response = self.__table.query(
            KeyConditionExpression=Key("username").eq(username)
        )
        items = response["Items"]

        # Return none, if the length of the items list is not 1
        if len(items) != 1:
            return None

        # Return the value fetched from the database
        return items[0]

    def delete_user_by_username(self, username: str):
        response = self.__table.delete_item(
            Key={"username": username}
        )

        return response
Enter fullscreen mode Exit fullscreen mode

Generally, I use double-underscore's in front of methods that will be used internally within a class and are not required outside of it. I feel this to be a good coding standard. However, you are open to use the way you feel like.

Testing

Moving to the testing of this service, create two files __init__.py and test_service.py files under the test directory. As with the src files, the __init__.py will be used for setting the current directory as module and will be empty. The test_service.py will be where our test code is going to be.

I am using a combination of pytest and unittest modules to perform testing.

Let's start by writing the setUp and tearDown methods. These methods are automatically whenever a test case is executed from the test suite.

from MockDynamoDbExample.src.service import Service
from moto import mock_dynamodb
from unittest import TestCase


@mock_dynamodb
class TestService(TestCase):
    def setUp(self) -> None:
        """Method called before every test case is executed"""

        # Create an instance of the Service class
        self.service = Service("user", env="test")

    def tearDown(self) -> None:
        """Method called after every test case has finished execution"""

        # Clear the table after use
        self.service.delete_table()

        # Reset the service variable
        self.service = None
Enter fullscreen mode Exit fullscreen mode

I am importing few things here, like the Service class that we implemented, the mock_dynamodb decorator from the moto library and TestCase from the unittest module.

The mock_dynamodb is going to be helpful in mocking all of our Dynamo Db related calls. This import will be used as a decorator. To learn more about decorators in general, follow this link.

In the setUp method, we are creating a Service class instance with "user" as the table name and environment as "test". The tearDown method implements the cleanup functionality. Both of these methods are called before and after every test case respectively.

Moving on, the following code implements the first test case which tests for create_user method in services

def test_create_user(self):
    username = "test-mike"
    age = 19
    first_name = "Mike"
    last_name = "Ross"
    email = "test_mike@example.com"

    try:
        self.service.create_user(username, email, age, first_name, last_name)
    except Exception as exc:
        assert False, f"service.create_user raised an exception {exc}"
    assert True
Enter fullscreen mode Exit fullscreen mode

We are testing the creation success using assert True at the end.

Let's run the test case by going over to the terminal or command prompt. Make sure that you are in the root directory of this project while executing the following command:

pytest test/test_service.py
Enter fullscreen mode Exit fullscreen mode

It should print out the following on Linux system

(venv) ibi@ibi:~/MockDynamoDbExample$ pytest MockDynamoDbExample/test/test_service.py 
======================================================================================== test session starts =========================================================================================
platform linux -- Python 3.9.12, pytest-7.1.1, pluggy-1.0.0
rootdir: /home/ibi/MockDynamoDbExample
collected 1 items                                                                                                                                                                                    

MockDynamoDbExample/test/test_service.py ...                                                                                                                                                   [100%]

========================================================================================= 1 passed in 1.46s ==========================================================================================
Enter fullscreen mode Exit fullscreen mode

For windows powershell, the output should look like

(venv) PS D:\MockDynamoDbExample> pytest test\test_service.py
================================================= test session starts =================================================
platform win32 -- Python 3.9.10, pytest-7.1.1, pluggy-1.0.0
rootdir: D:\MockDynamoDbExample
collected 1 items

test\test_service.py ...                                                                                         [100%]

================================================== 1 passed in 1.63s ==================================================
Enter fullscreen mode Exit fullscreen mode

Adding another test case to test getting a user by username.

def test_get_user_by_username(self):
    username = "test-mike"
    age = 19
    first_name = "Mike"
    last_name = "Ross"
    email = "test_mike@example.com"

    # Create a user in the database
    try:
        self.service.create_user(username, email, age, first_name, last_name)
    except Exception as exc:
        assert False, f"service.create_user raised an exception {exc}"

    # Get the user by username from the database
    try:
        response = self.service.get_user_by_username(username)
    except Exception as exc:
        assert False, f"service.get_user_by_username raised an exception {exc}"

    # Verify the details
    assert response is not None
    assert response["username"] == username
    assert response["age"] == age
    assert response["email"] == email
    assert response["first_name"] == first_name
    assert response["last_name"] == last_name
Enter fullscreen mode Exit fullscreen mode

In this test case, we are testing the insertion and fetch performed by Dynamo Db.

The complete file should look like

from MockDynamoDbExample.src.service import Service
from moto import mock_dynamodb
from unittest import TestCase


@mock_dynamodb
class TestService(TestCase):
    def setUp(self) -> None:
        """Method called before every test case is executed"""

        # Create an instance of the Service class
        self.service = Service("user", env="test")

    def tearDown(self) -> None:
        """Method called after every test case has finished execution"""

        # Clear the table after use
        self.service.delete_table()

        # Reset the service variable
        self.service = None

    def test_create_user(self):
        username = "test-mike"
        age = 19
        first_name = "Mike"
        last_name = "Ross"
        email = "test_mike@example.com"

        try:
            self.service.create_user(username, email, age, first_name, last_name)
        except Exception as exc:
            assert False, f"service.create_user raised an exception {exc}"
        assert True

    def test_get_user_by_username(self):
        username = "test-mike"
        age = 19
        first_name = "Mike"
        last_name = "Ross"
        email = "test_mike@example.com"

        # Create a user in the database
        try:
            self.service.create_user(username, email, age, first_name, last_name)
        except Exception as exc:
            assert False, f"service.create_user raised an exception {exc}"

        # Get the user by username from the database
        try:
            response = self.service.get_user_by_username(username)
        except Exception as exc:
            assert False, f"service.get_user_by_username raised an exception {exc}"

        # Verify the details
        assert response is not None
        assert response["username"] == username
        assert response["age"] == age
        assert response["email"] == email
        assert response["first_name"] == first_name
        assert response["last_name"] == last_name
Enter fullscreen mode Exit fullscreen mode

Additionally, if you don't want to use a class based approach and would like to write independent functions then you can do the following

@mock_dynamodb
def test_create_user():
    service = Service("user", "test")
    username = "test-mike"
    age = 19
    first_name = "Mike"
    last_name = "Ross"
    email = "test_mike@example.com"

    try:
        service.create_user(username, email, age, first_name, last_name)
    except Exception as exc:
        assert False, f"service.create_user raised an exception {exc}"
    assert True
    service.delete_table()
Enter fullscreen mode Exit fullscreen mode

As you can see, it can get a little tedious if you have a lot of test cases to write with a start and end common functionality.

Conclusion

moto is an amazing library to mock any AWS Dynamo DB based calls. You can check out the development on their GitHub page.

A trivia, prior to the moto library version 3.1.0, the mock decorator to be used was @mock_dynamodb2. Make sure if you are using the version prior to 3.1.0, use this decorator instead or update to the version I used in this blog.

Top comments (0)