Unit testing Lambda functions

by Adam Świątkowski
10 mins read
AWS Lambda Unit Testing

Talking about unit testing can cause quite a stir. This is because there are many approaches to this topic. In my post I will describe how I try to unit test my Lambda functions, but I do not think that there are no better approaches. I will use both AWS CDK as IAC (eng. Infrastructure as Code) and pytest as a testing framework for Python.

Testing pyramid
Testing pyramid

What are unit tests?

Unit testing is the process of testing the smallest pieces of code as possible. Some structures might be difficult to be tested, that’s why it’s recommended to always have in mind if the structure that was created will be testable or not. It’s also worth to understand what is TDD (eng. Test-driven-development) and why this approach might be helpful in the process of software creation.

Unit testing Lambda functions sample architecture
Unit testing Lambda functions sample architecture

Before we start unit testing our Lambda’s code

Before we start I’d like to describe a little, what kind of architecture will be used to demonstrate unit testing. Above, there’s diagram with the architecture I’ve created. We have there:

  • API Gateway, which receives the GET request from the user and responds to him,
  • Lambda function which validates the request and gets the parameter from the SSM Parameter Store
  • IAM Role which allows Lambda to access the SSM Parameter Store
Response of the API with parameter from SSM parameter store after providing valid user_id
Response of the API with parameter from SSM parameter store after providing valid user_id

Above infrastructure was built with AWS CDK and Python 3.9. Link to my repository with the source code is here.

Unit testing Lambda functions built with AWS CDK
Unit testing Lambda functions built with AWS CDK

How to properly organize Lambda function’s code

At the beginning of this post, I mentioned that it’s crucial to think about proper organizations of your code for it to be testable. There’s a great guide from AWS of how to properly decouple Lambda function’s code here. I’ll quote some of the most important rules in this post also.

  1. Separate business logic and Lambda handler – The handler should extract all required information from the event object and then call a separate method that implements the business logic.
  2. Minimize your deployment package size to its runtime requirements – Only include in Lambda’s package minimal dependencies needed for the function to properly work.
  3. Lambda should have idempotent code – Lambda should have the same behaviour no matter what happened in previous invocation.

Understanding the above will result in having testable, cost-optimized and optimal Lambda functions.

Keeping Lambda built with understanding best-practices will result in having cost-optimized and time-efficient product
Keeping Lambda built with understanding best-practices will result in having cost-optimized and time-efficient product

Let’s test something!

Below is the code of our Lambda function. As mentioned at the beginning, it validates user_id provided by the user in the request, and if it’s valid id and is longer than 3 characters Lambda connects to the SSM Parameter Store to get the parameter called DummySSM with value: My Real SSM Value.

Unit testing Lambda function
Unit testing Lambda function

Unit testing each of the Lambda’s part

First of all we need to understand what can be tested. Focus on our function’s behaviour and point how the function can behave:

  • Function can be properly called with proper user_id – We should receive HTTP STATUS 200 (OK),
  • Function can be properly called with invalid user_id (e.g. less character than 3) – We should receive HTTP STATUS 400 (Bad request),
  • Function can be called not properly or can’t connect to SSM – We should receive HTTP STATUS 500 (Internal Server Error).

Secondly, as we understand how the function can behave let’s think about ‘units’ to be tested. Units in our example will be functions handler and _prepare_response. Let’s start with _prepare_response function:

Unit testing Lambda business logic function 1
Unit testing Lambda business logic function 1

Our first scenario will be having the properly called function with proper input (user_id exists and is greater than 3 characters). The test should check if with the above assumptions function will return HTTP STATUS 200. Test example:

Unit testing Lambda business logic function 2
Unit testing Lambda business logic function 2

Initially, we create an empty Lambda Context and store in its parameters the name of the fake parameter from the SSM Parameter Store. After that we prepare request with proper, random user_id and after that we simply call the function which we want to test. Let’s run the above test.

pytest --log-cli-level=DEBUG -vv tests/unit
Pytest execution
Pytest execution

That’s it! As you see with prepared circumstances, function did its job properly. Exercise for you. Prepare test scenarios for HTTP STATUS 400. If you want to check if you did this well check my repository here.

How to check if Exception was raised with pytest?

Above all, you might ask how to test the situation when there won’t be a connection to SSM Parameter Store which in our case should result in raising an Exception. Anyway, I won’t keep you in suspense and will just tell you that there’s a raises() method of pytest. Method asserts that a block of code or a function call raises the specified exception.

pytest.raises() method for Unit Testing
pytest.raises() method for Unit Testing

Unit testing decorated Lambda function’s handler

Matrioszka / Babusha showing Decorator design pattern
Matrioszka / Babushka showing Decorator design pattern

By definition, a decorator is a function that takes another function and extends the behavior of the latter function without explicitly modifying it (just like Babushka/Matrioszka have another one inside). This fact can cause problems when testing a handler wrapped by a decorator.

Decorated Lambda function's handler
Decorated Lambda function’s handler

In our example we use @ssm_parameter_store decorator which connects to SSM Parameter Store inside our function to grab the parameter from the store. Unit testing should be focused on the code and logic behind it that’s why connection to other AWS services such as SSM Parameter Store should not be tested at this stage. This integration between Lambda and SSM service should be tested, as name suggests, during integration tests.

If we’ll now simply call our Lambda’s handler function we’ll encounter problem with credentials as the handler is decorated and @ssm_parameter_store function will try to get the parameter from the SSM.

ssm_parameter_store decorator logic
ssm_parameter_store decorator logic

To get around this problem, thus not connecting to the SSM Parameter Store and substituting the value returned by the decorator, we must mock the decorator and re-initialize our handler. Mocking is almost always required for unit tests because it provides a way to control how other modules behave.

MonkeyPatch() decorator for unit tests
MonkeyPatch() decorator for unit tests

Firstly, we patch the ssm_parameter_store decorator from the lambda_decorators package with mock_ssm_parameter_store_decorator function which will substitute original’s logic with our function with the below code:

Mocking decoreated Lambda Functions with pytest
Mocking decoreated Lambda Functions with pytest

As you see it’s more or less the same code as the original decorator ssm_parameter_store the change here is that we return fake value and don’t connect to the SSM Parameter Store. As a result, our Lambda’s handler will be ‘happy’ as it will receive all the things it needs.

Reimporting redecorated Lambda’s handler

Next, we need to reimport the handler’s content using importlib and return it as a redecorated_handler.

At end, we can build final test! Let’s test if our handler will result with HTTP STATUS CODE 400 when we’ll pass too short user_id.

Running redecorated Lambda function unit test
Running redecorated Lambda function unit test
Pytest execution result
Pytest execution result

That’s it! As you see we’ve managed to unit test our Lambda’s handler as well. Exercise for you! Prepare tests for handler returning code 200 and 500. If you want to check if you did this well check my repository here.

Different approaches with unit testing
Different approaches with unit testing

Different unit testing approaches

I know that the topic of unit testing can arouse various emotions. Please remember that mentioned methods are the ones I chose as a solution, I don’t consider it the only way of testing in the world. Especially in the topic of mocking AWS services, in more advanced cases, it is worth looking at the moto library, which can help in preparing responses from various AWS services without invoking them.

Summarize

Thank you for reaching out to that place. If you want to know more about AWS and the cloud, check the below posts:

Related Articles

This website uses cookies to improve your experience. I'll assume you're ok with this, but you can opt-out if you wish. Read Privacy Policy. Accept