How to test Python Lambda functions

Andrei Maksimov

Andrei Maksimov

More and more companies nowadays are adopting Serverless technologies. AWS Lambda is the essential building block of any Serverless cloud application or automation. Whatever AWS Lambda is doing for you, it is very important to test its code to make sure it’s working exactly as expected. In this article, we’ll cover manual and automated testing of Python Lambda functions that are interacting with DynamoDB and S3 services.

AWS Lambda function example

As an example let’s test a Lambda function that reads the file uploaded to the S3 bucket and saves its content to DynamoDB.

So, here are several places, which need to be tested:

Lambda function events
  • File upload to S3 – we need to make sure that during test cycle we’ll be dealing with the same file and the same content
  • File download – we need to make sure that our Lambda function can download, read and parse the file
  • Save the file to DynamoDB – we need to make sure, that our Lambda function code is able to save the file to DynamoDB of specified structure
  • Upload event – we need to make sure, that our Lambda function is able to process standard event from S3 bucket

Here’s the Lambda function code that we’re going to test:

'''
Simple Lambda function that reads file from S3 bucket and saves
its content to DynamoDB table
'''
#!/usr/bin/env python3

import json
import os
import logging
import boto3

LOGGER = logging.getLogger()
LOGGER.setLevel(logging.INFO)

AWS_REGION = os.environ.get('AWS_REGION', 'us-east-1')
DB_TABLE_NAME = os.environ.get('DB_TABLE_NAME', 'amaksimov-test')

S3_CLIENT = boto3.client('s3')
DYNAMODB_CLIENT = boto3.resource('dynamodb', region_name=AWS_REGION)
DYNAMODB_TABLE = DYNAMODB_CLIENT.Table(DB_TABLE_NAME)

def get_data_from_file(bucket, key):
    '''
    Function reads json file uploaded to S3 bucket
    '''
    response = S3_CLIENT.get_object(Bucket=bucket, Key=key)

    content = response['Body']
    data = json.loads(content.read())

    LOGGER.info('%s/%s file content: %s', bucket, key, data)

    return data

def save_data_to_db(data):
    '''
    Function saves data to DynamoDB table
    '''
    result = DYNAMODB_TABLE.put_item(Item=data)
    return result

def handler(event, context):
    '''
    Main Lambda function method
    '''
    LOGGER.info('Event structure: %s', event)

    for record in event['Records']:
        s3_bucket = record['s3']['bucket']['name']
        s3_file = record['s3']['object']['key']

        data = get_data_from_file(s3_bucket, s3_file)

        for item in data:
            save_data_to_db(item)

    return {
        'StatusCode': 200,
        'Message': 'SUCCESS'
    }

The code itself is self-explanatory, so we’ll not spend time describing it and move straight to the testing.

So, how can we make sure that the code above is actually working?

There are two ways of doing that:

  • Deploy and test AWS Lambda manually
  • Unit test AWS Lambda automatically before the deployment

Let’s start with the manual testing, first.

AWS Lambda manual testing

For manual testing of your AWS Lambda function you should have:

  • S3 bucket
  • DynamoDB table
  • Correct permissions for Lambda to access both

Let’s assume that you have all of them, if not, we recommend you to check our article Terraform – Deploy Lambda To Copy Files Between S3 Buckets to get mostly all Terraform deployment automation.

Example of test file / test data

Here’s the new_prices.json file content, which is uploaded to the prices folder at S3 bucket:

cat new_prices.json
[
    {"product": "Apple", "price": 15},
    {"product": "Orange", "price": 25}
]

This file is placed at S3 bucket to prices folder:

How to test Python Lambda functions - Test file location

Configuring test event

Now, to test the Lambda function manually, you need to open your Lambda function at the web console and select “Configure test event” from the dropdown menu of the “Test” button:

How to test Python Lambda functions - Configure test events

Now you can select a test event for almost any AWS service, which has an integration with Lambda. We need “Amazon S3 Put” example:

How to test Python Lambda functions - Amazon S3 Put

Next, you need to set up the test Event name, change the S3 bucket name, and test file location (key):

How to test Python Lambda functions - Configure S3 test event (final)

Finally, hit the “Create” button to save your test event.

Now, every single time you’ll click the “Test” button, it will send the same test event to your Lambda function.

If you did not deploy your Lambda function, do it now:

How to test Python Lambda functions - Deploy Lambda function

Lambda handler

Change Lambda function handler to lambda_function.handler:

How to test Python Lambda functions - Change Lambda handler

Lambda role (permissions)

If you’re doing manually the whole process, you need to add Lambda function required permissions to access DynamoDB and S3 bucket.

The quickest way to do this is to add them to the already created Lambda function role:

How to test Python Lambda functions - Configure permissions

For real production deployment, it is strongly recommended to provide only necessary permissions to your Lambda function, but for this example, we’ll add the following policies:

  • AmazonS3FullAccess
  • AmazonDynamoDBFullAccess
How to test Python Lambda functions - IAM policies

Test execution

Now, your Lambda has enough privileges, and you can press the “Test” button to test it manually:

How to test Python Lambda functions - Execution results

And in DynamoDB you should see your test data:

How to test Python Lambda functions - DynamoDB items

AWS Lambda test automation (unit-testing)

There are lots of downsides to manual operations. The most important are:

  • All manual operations are time consuming
  • Manual testing may introduce inconsistency and errors

Here’s where test automation can help.

For Python Lambda functions we can do it using the moto library. This library allows us to mock AWS services and test your code even before you deploy it. And it’s extremely fast!

Of course, not all the problems can be solved using moto, but if you’re using it, you can reduce serverless application development time up to 5 times (measurement is based on personal author experience).

Let’s take a look, how we can test the function above using moto.

Just to remind you, we’ll be testing those integrations:

Lambda function events

Requirements

For the examples below we’ll use the following Python libraries:

cat requirements_dev.txt
moto==2.0.2
coverage==5.5
pylint==2.7.2
pytest==6.2.2

You can install them in the virtual environment using the following commands:

python3 -m venv .venv
source .venv/bin/activate
pip install --upgrade pip
pip install -r requirements_dev.txt 

Project structure

In this article, we’ll use the following project structure:

tree
.
├── index.py
├── requirements_dev.txt
└── test_index.py

0 directories, 3 files

Where:

  • index.py is the file, which contain Lambda function code from the example below
  • requirements_dev.txt is the file which contains the required Python libraries for this article
  • test_index.py is the file, which contains test cases for the Lambda function code

Unit-testing AWS Lambda S3 file upload events

First, let’s check file upload to the S3 bucket and get_data_from_file() function:

import json
import unittest
import boto3
from moto import mock_s3

S3_BUCKET_NAME = 'amaksimov-s3-bucket'
DEFAULT_REGION = 'us-east-1'

S3_TEST_FILE_KEY = 'prices/new_prices.json'
S3_TEST_FILE_CONTENT = [
    {"product": "Apple", "price": 15},
    {"product": "Orange", "price": 25}
]

@mock_s3
class TestLambdaFunction(unittest.TestCase):

    def setUp(self):
        self.s3 = boto3.resource('s3', region_name=DEFAULT_REGION)
        self.s3_bucket = self.s3.create_bucket(Bucket=S3_BUCKET_NAME)
        self.s3_bucket.put_object(Key=S3_TEST_FILE_KEY,
                                  Body=json.dumps(S3_TEST_FILE_CONTENT))

    def test_get_data_from_file(self):
        from index import get_data_from_file

        file_content = get_data_from_file(S3_BUCKET_NAME, S3_TEST_FILE_KEY)
        self.assertEqual(file_content, S3_TEST_FILE_CONTENT)

In this simple test example, we’re importing required libraries, setting up global variables, and declaring TestLambdaFunction test case class, which has two methods

  • setUp() – this method is called before any executed test (function starts from test_ name) it is very useful when you need to reset the DB or other services to the initial state
  • test_get_data_from_file() – this test imports the get_data_from_file method, executes it and checks that returned file content is exactly what was uploaded (setUp()).

To execute this test, run the following command:

coverage run -m pytest

Having this test in place will allow us to make sure that we’re getting the same expected output from the get_data_from_file() method of our Lambda function.

Unit-testing AWS Lambda DynamoDB operations

Now, let’s add DynamoDB support to test save_data_to_db() and handler methods.

import json
import os
import unittest
import boto3
import mock
from moto import mock_s3
from moto import mock_dynamodb2

S3_BUCKET_NAME = 'amaksimov-s3-bucket'
DEFAULT_REGION = 'us-east-1'

S3_TEST_FILE_KEY = 'prices/new_prices.json'
S3_TEST_FILE_CONTENT = [
    {"product": "Apple", "price": 15},
    {"product": "Orange", "price": 25}
]

DYNAMODB_TABLE_NAME = 'amaksimov-dynamodb'

@mock_s3
@mock_dynamodb2
@mock.patch.dict(os.environ, {'DB_TABLE_NAME': DYNAMODB_TABLE_NAME})
class TestLambdaFunction(unittest.TestCase):

    def setUp(self):
        # S3 setup
        self.s3 = boto3.resource('s3', region_name=DEFAULT_REGION)
        self.s3_bucket = self.s3.create_bucket(Bucket=S3_BUCKET_NAME)
        self.s3_bucket.put_object(Key=S3_TEST_FILE_KEY,
                                  Body=json.dumps(S3_TEST_FILE_CONTENT))

        # DynamoDB setup
        self.dynamodb = boto3.client('dynamodb')
        try:
            self.table = self.dynamodb.create_table(
                TableName=DYNAMODB_TABLE_NAME,
                KeySchema=[
                    {'KeyType': 'HASH', 'AttributeName': 'product'}
                ],
                AttributeDefinitions=[
                    {'AttributeName': 'product', 'AttributeType': 'S'}
                ],
                ProvisionedThroughput={
                    'ReadCapacityUnits': 5,
                    'WriteCapacityUnits': 5
                }
            )
        except self.dynamodb.exceptions.ResourceInUseException:
            self.table = boto3.resource('dynamodb').Table(DYNAMODB_TABLE_NAME)

    def test_get_data_from_file(self):
        from index import get_data_from_file

        file_content = get_data_from_file(S3_BUCKET_NAME, S3_TEST_FILE_KEY)
        self.assertEqual(file_content, S3_TEST_FILE_CONTENT)

In addition to the previous code we’ve added DynamoDB table declaration with HASH key product of type string (S).

@mock.patch.dict sets up DB_TABLE_NAME environment variable used by the Lambda function.

Now, let’s test save_data_to_db() method:

def test_save_data_to_db(self):
    from index import save_data_to_db

    for item in S3_TEST_FILE_CONTENT:
        save_data_to_db(item)

    db_response = self.table.scan(Limit=1)

    db_records = db_response['Items']

    while 'LastEvaluatedKey' in db_response:
        db_response = self.table.scan(
            Limit=1,
            ExclusiveStartKey=db_response['LastEvaluatedKey']
        )
        db_records += db_response['Items']

    self.assertEqual(len(S3_TEST_FILE_CONTENT), len(db_records))

In this test, we’re calling the save_data_to_db() method for every single item from S3_TEST_FILE_CONTENT list to save it to DynamoDB.

Then we’re querying DynamoDB and checking that we’re getting back the same amount of records.

Finally, we can test the handler method.

This method will get the following S3 integration data structure in the event variable during its execution.

You may also copy-paste the event structure which we created during the manual test, but as for me, my event example is a bit more clear.

def test_handler(self):
        from index import handler

        event = {
            'Records': [
                {
                    's3': {
                        'bucket': {
                            'name': S3_BUCKET_NAME
                        },
                        'object': {
                            'key': S3_TEST_FILE_KEY
                        }
                    }
                }
            ]
        }

        result = handler(event, {})
        self.assertEqual(result, {'StatusCode': 200, 'Message': 'SUCCESS'})

Final test_index.py content is the following:

import json
import os
import unittest
import boto3
import mock
from moto import mock_s3
from moto import mock_dynamodb2

S3_BUCKET_NAME = 'amaksimov-s3-bucket'
DEFAULT_REGION = 'us-east-1'

S3_TEST_FILE_KEY = 'prices/new_prices.json'
S3_TEST_FILE_CONTENT = [
    {"product": "Apple", "price": 15},
    {"product": "Orange", "price": 25}
]

DYNAMODB_TABLE_NAME = 'amaksimov-dynamodb'

@mock_s3
@mock_dynamodb2
@mock.patch.dict(os.environ, {'DB_TABLE_NAME': DYNAMODB_TABLE_NAME})
class TestLambdaFunction(unittest.TestCase):

    def setUp(self):
        # S3 setup
        self.s3 = boto3.resource('s3', region_name=DEFAULT_REGION)
        self.s3_bucket = self.s3.create_bucket(Bucket=S3_BUCKET_NAME)
        self.s3_bucket.put_object(Key=S3_TEST_FILE_KEY,
                                  Body=json.dumps(S3_TEST_FILE_CONTENT))

        # DynamoDB setup
        self.dynamodb = boto3.client('dynamodb')
        try:
            self.table = self.dynamodb.create_table(
                TableName=DYNAMODB_TABLE_NAME,
                KeySchema=[
                    {'KeyType': 'HASH', 'AttributeName': 'product'}
                ],
                AttributeDefinitions=[
                    {'AttributeName': 'product', 'AttributeType': 'S'}
                ],
                ProvisionedThroughput={
                    'ReadCapacityUnits': 5,
                    'WriteCapacityUnits': 5
                }
            )
        except self.dynamodb.exceptions.ResourceInUseException:
            self.table = boto3.resource('dynamodb').Table(DYNAMODB_TABLE_NAME)

    def test_get_data_from_file(self):
        from index import get_data_from_file

        file_content = get_data_from_file(S3_BUCKET_NAME, S3_TEST_FILE_KEY)
        self.assertEqual(file_content, S3_TEST_FILE_CONTENT)

    def test_save_data_to_db(self):
        from index import save_data_to_db

        for item in S3_TEST_FILE_CONTENT:
            save_data_to_db(item)

        db_response = self.table.scan(Limit=1)

        db_records = db_response['Items']

        while 'LastEvaluatedKey' in db_response:
            db_response = self.table.scan(
                Limit=1,
                ExclusiveStartKey=db_response['LastEvaluatedKey']
            )
            db_records += db_response['Items']

        self.assertEqual(len(S3_TEST_FILE_CONTENT), len(db_records))

    def test_handler(self):
        from index import handler

        event = {
            'Records': [
                {
                    's3': {
                        'bucket': {
                            'name': S3_BUCKET_NAME
                        },
                        'object': {
                            'key': S3_TEST_FILE_KEY
                        }
                    }
                }
            ]
        }

        result = handler(event, {})
        self.assertEqual(result, {'StatusCode': 200, 'Message': 'SUCCESS'})

Running automated unit-tests

If you’ll run this test again, you may see that all 3 tests successfully passed:

How to test Python Lambda functions - Coverage run pytest

To get test execution statistics, execute the following coverage commands:

coverage html --omit='.venv/*'
open htmlcov/index.html

Here’s an expected web page:

How to test Python Lambda functions - Coverage report

Which can give you more information, what’s covered in your code by tests, and what is not:

How to test Python Lambda functions - Coverage report extended

Summary

In this article, we covered the process of manual and automatic unit testing of your AWS Lambda function code by testing the S3 file upload event and DynamoDB put_item and scan operations.

We hope you’ll find this article useful. If so, please, help us to spread it to the world!

How useful was this post?

Click on a star to rate it!

As you found this post useful...

Follow us on social media!

We are sorry that this post was not useful for you!

Let us improve this post!

Tell us how we can improve this post?

Related Udemy Courses

Subscribe to our updates

Like this article?

Share on facebook
Share on Facebook
Share on twitter
Share on Twitter
Share on linkedin
Share on Linkdin
Share on pinterest
Share on Pinterest

Want to be an author of another post?

We’re looking for skilled technical authors for our blog!

Leave a comment

If you’d like to ask a question about the code or piece of configuration, feel free to use https://codeshare.io/ or a similar tool as Facebook comments are breaking code formatting.