How to perform unit testing for AWS Lambda functions

Performing unit testing for Lambda functions follows similar principles to traditional unit testing for the most part. The key difference is that Lambda functions often interact with external services or resources, but the basic functionality of unit testing is to test the code as a standalone unit. Also, we don’t want the unit tests to make any changes to any of our resources. For example, if our Lambda function adds data to a database, we don’t want it to add real data while testing. That’s where mocking libraries come in. They enable us to simulate these resources and test our Lambda function code without making any actual changes in our infrastructure.

Perform unit testing for a Lambda function

Unit testing a Lambda function essentially means unit testing the code that will be executed within the Lambda function. Here are the steps to perform the unit testing for a Lambda function:

Step 1: Identify the testable components

We start by identifying the individual functions or modules that we want to test within our Lambda function. We break down our code into small, testable units to ensure isolated testing.

Given below is a Python code for a Lambda function that we want to test before deployment.

import json
import boto3
from botocore.exceptions import ClientError
dynamodb = boto3.resource('dynamodb')
table_name = 'ExampleTable'
def lambda_handler(event, context):
    operation = event.get('operation')
    table = dynamodb.Table(table_name)
    if operation == 'create':
        item = event.get('item')
        if not item:
            return {
                'statusCode': 400,
                'body': json.dumps({'error': 'Item data is required for create operation'})
            }
        table.put_item(Item=item)
        return {
            'statusCode': 200,
            'body': json.dumps({'message': 'Item created successfully'})
        }
    elif operation == 'read':
        item_id = event.get('id')
        if not item_id:
            return {
                'statusCode': 400,
                'body': json.dumps({'error': 'Item ID is required for read operation'})
            }
        try:
            response = table.get_item(Key={'id': item_id})
        except ClientError as e:
            return {
                'statusCode': 500,
                'body': json.dumps({'error': str(e)})
            }
        item = response.get('Item')
        return {
            'statusCode': 200,
            'body': json.dumps(item if item else {})
        }
    elif operation == 'update':
        item_id = event.get('id')
        item_data = event.get('item')
        if not item_id or not item_data:
            return {
                'statusCode': 400,
                'body': json.dumps({'error': 'Item ID and item data are required for update operation'})
            }
        table.update_item(
            Key={'id': item_id},
            UpdateExpression="set info=:info",
            ExpressionAttributeValues={
                ':info': item_data['info'],
            },
            ReturnValues="UPDATED_NEW"
        )
        return {
            'statusCode': 200,
            'body': json.dumps({'message': 'Item updated successfully'})
        }
    elif operation == 'delete':
        item_id = event.get('id')
        if not item_id:
            return {
                'statusCode': 400,
                'body': json.dumps({'error': 'Item ID is required for delete operation'})
            }
        table.delete_item(Key={'id': item_id})
        return {
            'statusCode': 200,
            'body': json.dumps({'message': 'Item deleted successfully'})
        }
    else:
        return {
            'statusCode': 400,
            'body': json.dumps({'error': 'Invalid operation'})
        }

We can implement unit tests to test these each of these operation of the Lambda function. Along with it, we can also test the response generation of the Lambda function to make sure that the response is in the required format.

So in total, we've identified five parts of the Lambda function code to be testable.

Step 2: Choose a testing framework

We select a testing framework that supports the programming language or technology stack we’re using for our Lambda functions. Our Lambda function code is written in Python. There are many testing libraries for python some of which are unittest (built-in testing library), pytest and doctest. We'll use pytest to test our Lambda function code.

Step 3: Write test cases

We create test cases for each unit of our Lambda functions. Test cases should cover both expected and edge cases to validate the functions’ behavior under different scenarios. We'll keep things simple and create just one test for each of the identified testable component of our code.

The test codes are given below:

import json
from lambda_function import lambda_handler  # Importing the lambda_handler function from lambda_function module
import boto3  # Importing the Boto3 library for AWS SDK
# Initialize a DynamoDB resource
dynamodb = boto3.resource('dynamodb')
# Name of the DynamoDB table
table_name = 'ExampleTable'
def test_create_item():
    # Define the event with operation 'create' and item details
    event = {
        'operation': 'create',
        'item': {'id': '123', 'info': 'test info'}
    }
    
    # Call the lambda_handler function with the event
    response = lambda_handler(event, None)
    
    # Assertions to verify the response
    assert response['statusCode'] == 200
    assert json.loads(response['body']) == {'message': 'Item created successfully'}
    # Verify the item was created in DynamoDB
    table = dynamodb.Table(table_name)
    item = table.get_item(Key={'id': '123'}).get('Item')
    assert item == {'id': '123', 'info': 'test info'}
if __name__ == "__main__":
    # Run the test function
    test_create_item()

Create item

In these test files, we are invoking the functions of our code with pre-defined values and checking if the response is as desired or not. If the response is as expected, the test passes or else it fails.

Step 4: Mock external dependencies

In our test codes, we're using the actual DynamoDB table that the Lambda is supposed to interact with. This causes two types of issues:

Our unit testing is not standalone because of this as any issue with DynamoDB service or this table might lead the test to fail.
We are making changes to our production environment during testing, which is not preferable.

To mitigate this issue, we'll use a mock library to simulate the DynamoDB table. Mocking helps ensure that the focus of the test is solely on the unit being tested and no unintentional change is made to the existing infrastructure. We'll use moto as our mock library. It allows us to simulate AWS resources, enabling us to execute testing without worrying about any unintentional changes to our AWS infrastructure.

Our test files will now look as follows:

import json
import boto3
from moto import mock_aws  # Importing moto for mocking AWS services
from lambda_function import lambda_handler  # Import the lambda_handler function from lambda_function module
@mock_aws
def test_create_item_response():
    dynamodb = boto3.resource('dynamodb')
    table_name = 'ExampleTable'
    # Create a mock DynamoDB table
    table = dynamodb.create_table(
        TableName=table_name,
        KeySchema=[{'AttributeName': 'id', 'KeyType': 'HASH'}],
        AttributeDefinitions=[{'AttributeName': 'id', 'AttributeType': 'S'}],
        ProvisionedThroughput={'ReadCapacityUnits': 5, 'WriteCapacityUnits': 5}
    )
    table.meta.client.get_waiter('table_exists').wait(TableName=table_name)  # Wait until the table exists
    # Define the event for creating an item
    event = {
        'operation': 'create',
        'item': {'id': '123', 'info': 'test info'}
    }
    # Call the lambda_handler function with the event
    response = lambda_handler(event, None)
    
    # Assertions to verify the response for creating an item
    assert response['statusCode'] == 200
    assert json.loads(response['body']) == {'message': 'Item created successfully'}
# Run the test function
if __name__ == "__main__":
    test_create_item_response()

Create item

We've used moto to simulate a DynamoDB table in these tests. The Lambda function code then interacts with this simulated table during testing. To do that, all we had to do was import mock_aws from moto and then wrap the test with it. It mocked out all the AWS calls. The last test does not require any interaction with external resources so we've skipped the table simulation for it.

Step 5: Execute tests

We run our unit tests using the chosen testing framework. The framework will execute the tests and report any failures or errors. We make sure to run the tests in an isolated environment to avoid interference from other components or dependencies.

We can simply test the Lambda function code locally and do not need to invoke the Lambda function to check the desired response. As discussed earlier, we'll use the pytest library to execute these tests. Click the "Run" button to open the terminal.

Note: The following environments have been set in the coding playground below as these are required to be set when using the moto library:
AWS_ACCESS_KEY_ID = testing
AWS_SECRET_ACCESS_KEY = testing
AWS_DEFAULT_REGION = us-east-1
As these are dummy variables, they can be set to any value but they must be set either within the code or within the execution environment.

import json
import boto3
from botocore.exceptions import ClientError

dynamodb = boto3.resource('dynamodb')
table_name = 'ExampleTable'

def lambda_handler(event, context):
    operation = event.get('operation')
    table = dynamodb.Table(table_name)

    if operation == 'create':
        item = event.get('item')
        if not item:
            return {
                'statusCode': 400,
                'body': json.dumps({'error': 'Item data is required for create operation'})
            }
        table.put_item(Item=item)
        return {
            'statusCode': 200,
            'body': json.dumps({'message': 'Item created successfully'})
        }

    elif operation == 'read':
        item_id = event.get('id')
        if not item_id:
            return {
                'statusCode': 400,
                'body': json.dumps({'error': 'Item ID is required for read operation'})
            }
        try:
            response = table.get_item(Key={'id': item_id})
        except ClientError as e:
            return {
                'statusCode': 500,
                'body': json.dumps({'error': str(e)})
            }
        item = response.get('Item')
        return {
            'statusCode': 200,
            'body': json.dumps(item if item else {})
        }

    elif operation == 'update':
        item_id = event.get('id')
        item_data = event.get('item')
        if not item_id or not item_data:
            return {
                'statusCode': 400,
                'body': json.dumps({'error': 'Item ID and item data are required for update operation'})
            }
        table.update_item(
            Key={'id': item_id},
            UpdateExpression="set info=:info",
            ExpressionAttributeValues={
                ':info': item_data['info'],
            },
            ReturnValues="UPDATED_NEW"
        )
        return {
            'statusCode': 200,
            'body': json.dumps({'message': 'Item updated successfully'})
        }

    elif operation == 'delete':
        item_id = event.get('id')
        if not item_id:
            return {
                'statusCode': 400,
                'body': json.dumps({'error': 'Item ID is required for delete operation'})
            }
        table.delete_item(Key={'id': item_id})
        return {
            'statusCode': 200,
            'body': json.dumps({'message': 'Item deleted successfully'})
        }

    else:
        return {
            'statusCode': 400,
            'body': json.dumps({'error': 'Invalid operation'})
        }

Test Lambda function code

All our tests will execute and will pass because our code is functioning as required. Now here's a challenge for you—try and modify the Lambda function code that makes one of the tests fail and then try to modify the test accordingly.

So we've seen how we can implement unit testing for Lambda functions to ensure their reliability and functionality. This testing is important and itensures that the code behaves as expected under various scenarios, maintaining the integrity of the serverless application.

Free Resources