CODEWITHSUNDEEP
pythonunit-testingaws-lambdapython-unittest
kodi1911
kodi1911

Reputation: 722

Unit test for Python AWS Lambda: mock function before module is imported

I'm trying to write unit tests for my aws lambda function written in python 3.9. I tried different things to mock the get_object function that makes calls to the S3. I wanted to focus only on the calculate method, to verify if I'm getting correct results of an calculation.

When I try to run the following approach I'm getting credential errors about boto3

python -m unittest tests/app-test.py
...
botocore.exceptions.NoCredentialsError: Unable to locate credentials

Is there a way to import the calculate method from app.py and mock call to the get_object fn?

directory:

functions:
- __init__.py
- app.py
tests:
- __init__.py
- app-test.py

lambda function app.py:

import json
import boto3

def get_object():
    s3 = boto3.client('s3')
    response = s3.get_object(Bucket='mybucket', Key='object.json')
    content = response['Body'].read().decode('utf-8')
    return json.loads(content)


stops = get_object()


def lambda_handler(event, context): 
    params = event['queryStringParameters']
    a = int(params['a'])
    b = int(params['b'])
   
    result = calculate(a, b)
    return {
        'statusCode': 200,
        'body': json.dumps(result)
    }


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

unit test app-test.py:

import unittest
from unittest import mock

with mock.patch('functions.app.get_object', return_value={}):
    from functions.app import calculate

class TestCalculation(unittest.TestCase):
   def test_should_return_correct_calculation(self):
     # when
     result = calculate(1, 2)

     # then
     self.assertEqual(3, result)

Upvotes: 2

Views: 3874

Answers (2)

Alex
Alex

Reputation: 173

Actually a simpler approach is to set the environment variables AWS_ACCESS_KEY_ID and AWS_SECRET_ACCESS_KEY before the import of the components that use AWS SDK (boto3 etc.)

import os

os.environ['AWS_ACCESS_KEY_ID'] = "ABC"
os.environ['AWS_SECRET_ACCESS_KEY'] = "123"

from module_that_uses_boto3 import ...

Upvotes: 0

kodi1911
kodi1911

Reputation: 722

I was able to fix the issue. The biggest obstacle was to mock the boto3 in the app.py. I did this, by mocking the whole boto3 module before it's imported. Here's the code of app-test.py

import sys
from io import BytesIO
from json import dumps
from unittest import TestCase, main
from unittest.mock import Mock
from botocore.stub import Stubber
from botocore.session import get_session
from botocore.response import StreamingBody


# prepare mocks for boto3
stubbed_client = get_session().create_client('s3')
stubber = Stubber(stubbed_client)

# mock response from S3
body_encoded = dumps({'name': 'hello world'}).encode()
body = StreamingBody(BytesIO(body_encoded), len(body_encoded))
stubbed.add_response('get_object', {'Body': body})

stubber.activate()

# add mocks to the real module
sys.modules['boto3'] = Mock()
sys.modules['boto3'].client = Mock(return_value=stubbed_client)


# Import the module that will be tested
# boto3 should be mocked in the app.py
from functions.app import calculate


class TestCalculation(TestCase):
   def test_should_return_correct_calculation(self):
     # when
     result = calculate(1, 2)

     # then
     self.assertEqual(3, result)

Upvotes: 1

Related Questions