Mocking AWS Dynamo Db calls during testing
Tags: dynamodb, moto, testing, python
Published: 2022-04-24 03:16:24

## 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.
<br />

## 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](https://docs.python.org/3.9/library/venv.html) 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.
```
Installs pytest version 7.1.1

```
pip install moto[all]==3.1.5
```
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
```
Installs boto3 version 1.21.45


You can check out the complete list of requirements for this project on my [GitHub](https://github.com/iamibi/Mock-AWS-Services/tree/MockDynamoDbExample/) page.
<br />

## 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
```
<br>

## 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)
```
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()
```
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
```
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](https://boto3.amazonaws.com/v1/documentation/api/latest/guide/dynamodb.html) 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
```

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.
<br>

## 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
```
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](https://realpython.com/primer-on-python-decorators/) 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
```
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
```

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 ==========================================================================================
```

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 ==================================================
```

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
```

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
```

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()
```
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.
<br>

## Conclusion
**moto** is an amazing library to mock any AWS Dynamo DB based calls. You can check out the development on their [GitHub](https://github.com/spulec/moto) 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.
<br>