Python unittest example for DynamoDB operations using moto library

MoviesCreateTable.py

from pprint import pprint
import boto3

def create_movie_table(dynamodb=None):
    if not dynamodb:
        dynamodb = boto3.resource('dynamodb', endpoint_url='http://localhost:8000')

    table = dynamodb.create_table(
        TableName='Movies',
        KeySchema=[
            {
                'AttributeName': 'year',
                'KeyType': 'HASH'
            },
            {
                'AttributeName': 'title',
                'KeyType': 'RANGE'
            }
        ],
        AttributeDefinitions=[
            {
                'AttributeName': 'year',
                'AttributeType': 'N'
            },
            {
                'AttributeName': 'title',
                'AttributeType': 'S'
            }
        ],
        ProvisionedThroughput={
            'ReadCapacityUnits': 1,
            'WriteCapacityUnits': 1
        }
    )

    # Wait until the table exists.
    table.meta.client.get_waiter('table_exists').wait(TableName='Movies')
    assert table.table_status == 'ACTIVE'

    return table

if __name__ == '__main__':
    movie_table = create_movie_table()
    print("Table status:", movie_table.table_status)

MoviesGetItem.py

from pprint import pprint
import boto3
from botocore.exceptions import ClientError


def get_movie(title, year, dynamodb=None):
    if not dynamodb:
        dynamodb = boto3.resource('dynamodb', endpoint_url="http://localhost:8000")

    table = dynamodb.Table('Movies')

    try:
        response = table.get_item(Key={'year': year, 'title': title})
    except ClientError as e:
        print(e.response['Error']['Message'])
    else:
        return response['Item']


if __name__ == '__main__':
    movie = get_movie("The Big New Movie", 2015)
    if movie:
        print("Get movie succeeded:")
        pprint(movie, sort_dicts=False)

MoviesPutItem.py

from pprint import pprint
import boto3
from botocore.exceptions import ClientError

def put_movie(title, year, plot, rating, dynamodb=None):
    if not dynamodb:
        dynamodb = boto3.resource('dynamodb', endpoint_url="http://localhost:8000")

    table = dynamodb.Table('Movies')
    response = table.put_item(
       Item={
            'year': year,
            'title': title,
            'info': {
                'plot': plot,
                'rating': rating
            }
        }
    )
    return response

if __name__ == '__main__':
    movie_resp = put_movie("The Big New Movie", 2015,
                           "Nothing happens at all.", 0)
    print("Put movie succeeded:")
    pprint(movie_resp, sort_dicts=False)

TestMovies.py

from pprint import pprint
import unittest
import boto3
from botocore.exceptions import ClientError
from moto import mock_dynamodb2

@mock_dynamodb2
class TestDatabaseFunctions(unittest.TestCase):
    def setUp(self):
        """Create the mock database and table"""
        self.dynamodb = boto3.resource('dynamodb', region_name='us-east-1')
        
        from MoviesCreateTable import create_movie_table
        self.table = create_movie_table(self.dynamodb) 
    
    def tearDown(self):
        """Delete mock database and table after test is run"""
        self.table.delete()
        self.dynamodb=None

    def test_table_exists(self):
        self.assertTrue(self.table) # check if we got a result
        self.assertIn('Movies', self.table.name) # check if the table name is 'Movies'
        # pprint(self.table.name)

    def test_put_movie(self):
        from MoviesPutItem import put_movie

        result = put_movie("The Big New Movie", 2015,
                           "Nothing happens at all.", 0, self.dynamodb)
        
        self.assertEqual(200, result['ResponseMetadata']['HTTPStatusCode'])
        # pprint(result, sort_dicts=False)

    def test_get_movie(self):
        from MoviesPutItem import put_movie
        from MoviesGetItem import get_movie

        put_movie("The Big New Movie", 2015,
                    "Nothing happens at all.", 0, self.dynamodb)
        result = get_movie("The Big New Movie", 2015, self.dynamodb)

        self.assertEqual(2015, result['year'])
        self.assertEqual("The Big New Movie", result['title'])
        self.assertEqual("Nothing happens at all.", result['info']['plot'])
        self.assertEqual(0, result['info']['rating'])
        # pprint(result, sort_dicts=False)

if __name__ == '__main__':
    unittest.main()