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.
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.
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
Above infrastructure was built with AWS CDK and Python 3.9. Link to my repository with the source code is here.
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.
- 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.
- Minimize your deployment package size to its runtime requirements – Only include in Lambda’s package minimal dependencies needed for the function to properly work.
- 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.
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 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:
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:
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
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.
Unit testing decorated Lambda function’s handler
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.
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.
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.
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:
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.
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 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.
Thank you for reaching out to that place. If you want to know more about AWS and the cloud, check the below posts:
- Air pollution monitoring system built with AWS IoT and Greengrass
- How to run your own mail server on Amazon AWS EC2
- Is AWS certification really worth it?