How to write Unit Test for AWS Lambda Function in Python

In this tutorial, you will learn how to write unit test for AWS Lambda Function in Python.

Unit testing is a software testing method in which a specific piece of code is tested and make sure that the code is working as expected. Unit testing allows any faults or errors in the code to be found early and corrected before the final product is delivered.

In unit testing, the interaction of the code with another part of the application or external systems is frequently mocked.

The test functionality in the Lambda console is useful for testing small projects. This may involve deploying the function code, creating a test event, pressing the Test button, and verifying the result. However, if our project becomes large enough, the Lambda inline editor will not be shown and we will be unable to edit the function code directly in the Lambda console. A minor change to the function code necessitates a thorough test. As a result, unit testing can be useful in ensuring that the function code is working without any bug or error. We can also set up a CI/CD pipeline to perform tests on every deployment. This is yet another useful check that may prevent us from deploying code that does not work.

Let the code of your Lambda function be the following:

import json

def lambda_handler(event, context):
    print(event)
    numbers = event['numbers']
    response = json.dumps({"addition":sum(numbers[0],numbers[1]), "multiplication":multiply(numbers[0],numbers[1])})
    return {
        "statusCode": 200,
        "body": response,
    }

def sum(a,b):
    return a+b

def multiply(a,b):
    return a * b

Create a new test file named test_handler.py with the following code:

import json
import unittest
from src import lambda_function

class calculationTest(unittest.TestCase):

    def setUp(self):
        self.event = {"numbers":[5,7]}
        self.a = 5
        self.b = 7

    def test_lambda_handler(self):
        result = lambda_function.lambda_handler(self.event,'')
        data = json.loads(result["body"])
        print("data: ",data)
        expected_response = {"addition": 12, "multiplication": 35}
        self.assertEqual(data, expected_response)

    def test_sum(self):
       result = lambda_function.sum(self.a,self.b)
       self.assertEqual(result, self.a + self.b)


    def test_func_multiply(self):
        result = lambda_function.multiply(self.a,self.b)
        self.assertEqual(result, self.a * self.b)

To test, open a terminal and navigate to the src folder of your Lambda project and run the following command:

python -m unittest discover

The above test should give the output as follows:

----------------------------------------------------------------------
Ran 3 tests in 0.000s

OK

Our tests passed!