Writing unit tests for Lambda functions in Python

Code Jun 28 2021

This post shares what I've learned about writing unit tests for Lambda functions. I'll explain what unit tests are and why they can help you write and make changes to your function code quickly. I've also written an example Python function and unit test (using the unittest framework) so you can see it in practice.

As a relatively new programmer, writing tests is a new concept for me. Over the past year I've been adding unit tests into my largest side project (Haohaotiantian) and now I write them for any new project I start. I've found unit tests helpful because they allow me to quickly check the basic functionality of my code without needing to wait for a full deployment to the cloud. I've also integrated my tests into my CI/CD pipeline to test the code on every deploy and stop my deployment if the code is broken.

There are a variety of ways to test Lambda functions and serverless applications - like unit testing and integration testing, and testing locally or in the cloud. This post will show you one way to unit test Lambda function code locally that's worked well for me!

Jump to:

What are unit tests

Unit tests are written to test a specific piece of code and make sure it operates in the way you want it to (More details on different types of testing here). A unit test will ensure that your code works and that the inputs and outputs of the code are as expected.

In order to focus on just one piece of code, unit tests will often mock the code's interaction with other parts of the application or external systems. For example, a unit test for your Lambda function might include a fake response similar to what you expect to receive from the DynamoDB API call that it makes. Because you don't need to actually interact with AWS service APIs in the unit test, you can run unit tests locally on your computer without needing to deploy your function to the cloud.

Why are unit tests useful when writing Lambda function code

Using the test functionality in the Lambda console (deploying function code, creating a test event, and pressing 'Test') worked well for me for my first projects and is a good way to get started. However, deploying a new version of your application can take a few minutes, which is a while to wait if you just want to test a small code change. If your function code package is larger than 3 MB, you will also run into the limitation of not being able to edit your function code in the console. This means that even a small iteration to the function code requires a full deploy to test it.

This is where unit tests are helpful! While you're writing or making changes to your code, you can run your tests locally from within your text editor. This will help you catch errors, like syntax issues or misspelled variable names, without needing to wait for a deployment. You can also configure your CI/CD pipeline to run your unit tests on every new deploy. This is another helpful check that will prevent you deploying code that doesn't work.

Walkthrough of an example Python function and unit test

You can clone this example repo and try out the test for yourself: lambda-python-unit-test-example

SAM template & function code

Let's review the app. This app takes a text file to be translated and outputs the translation and some metadata.

Architecture diagram

I have a SAM template that defines an S3 bucket and a Lambda function. The translate_file Lambda function is triggered whenever a new file is uploaded to the S3 bucket.

The function code gets the S3 bucket and key for the file, checks if it's a valid file type (.txt extension), reads the file, and then calls the AWS Translate API to translate the file text. I've included two example files if you want to deploy the app and test the functionality.


Unit test

Now let's check out the unit test.

Since my function code is in a different directory than my test, I've updated the system path to be able to access the function code. Next I import the Python unittest framework and my function handler.


Mocking API calls

The next two functions are mocking my read_file and translate_text functions because these make calls to AWS services. The mock functions return what I expect to get back from these functions (based on the input event at the bottom of my unit test, which I'll get to in a moment). This allows me to run the test locally without actually calling any APIs. If you need to mock an integration with another system (like an API call), I find it helpful to isolate that call in it's own function so that you can mock the response of the wrapper function.


Thinking about what to test

Next up, I've created my test class, and within that my two main test functions, test_valid_file and test_invalid_file. I've added decoraters (@mock.patch) to each test function to swap in my mocked functions in place of the real ones.

When writing tests, we need to think about what the function inputs and outputs will look like when everything goes right. We also need to think about all the types of incorrect inputs our funcion might get or where it might fail, and what type of error messages we want it to output when it does.

The two main cases I want to check for are how my code handles valid (.txt) and invalid (anything else) file types. There are plenty of other scenarios I could potentially test (like if my call to get the S3 file or translate the text returns an error), but I'm keeping it simple to make it easier to understand the core components. How comprehensive you want your tests to be is up to you!

The rest of the main test functions are assertions. These tell the test what we expect the function to output, how many times we expect to call the mock functions, and more. Here's a full list of types of assertions. If any of the assertions fail, the test will fail and we know our code isn't working as we expect it to.


Test event payloads

Within each of my test functions, I set the file name (with a valid or invalid file extension). I pass the file name to my test event (at the bottom of the test, s3_upload_event). My Lambda handler invokes itself with the contents of s3_upload_event.


There are many ways you can invoke a Lambda function, and it's important to know what those invocation events will look like. You can find example event payloads by in service documentation, or you can do a test invoke and print out the event contents within your function. I run the printed event contents through a JSON validator to quickly add indentation and make the event more readable.

If you're going to publish your code as open source, be sure to replace any personal data or ids that might be in the example event!

Try it out

Now that we've walked through the code and unit test, let's try it out!

Make sure you have a default AWS region configured, since Python SDK boto3 requires it to create a service client. I use the AWS Toolkit for VS Code to configure my credentials and region in my text editor.

Open a terminal in your text editor and navigate to the src/tests/unit directory. Once there, run the unittest command,

python3 -m unittest discover

It should output a response like this:

Test results, OK

Our tests passed!

Now let's break the tests. Go into the function and make some changes, maybe comment out extracting the bucket_name and key_name from the event or change the invalid file type response. Remember that we're mocking the read_file and translate_text functions, so any changes in those won't be picked up by the test. Now when you run the same command, the tests will fail.

Test results, FAILED

Add tests to your CI/CD pipeline

Once you've written a unit test, the next cool thing you can do with it is add it as a step in your CI/CD pipeline. I've written more in depth about setting up a CI/CD pipeline for your serverless apps in this post, so check it out if that's new to you!

I've added an example .circleci/config.yml file in the sample code so you can see how to add tests to your CI/CD configuration. I've just added a 'Run tests' step, and now every time the app is deployed, it will run the tests. If the tests fail, the deployment will also fail. This will help you avoid deploying broken code.


Thanks for reading! Hopefully testing your function code locally will help you write code more quickly and build faster.

🌻

← Back Home