Unit testing with Python

Like most developers who started using AWS technology, especially Lambda functions, I found it necessary to start writing code in Python, and with that switch from Java to Python, I started looking for libraries and techniques to test my code.

Most of the developers including myself find it hard to test our code when we face some complicated scenarios, I will be walking you through some of these cases to show you how to solve these cases, so let’s rock and roll.

Pasadena, California

Case I: How to mock datetime?

FreezGun library is a library that will help you to freeze the time:

Once the decorator or context manager have been invoked, all calls to datetime.datetime.now(), datetime.datetime.utcnow(), datetime.date.today(), time.time(), time.localtime(), time.gmtime(), and time.strftime() will return the time that has been frozen. time.monotonic() and time.perf_counter() will also be frozen, but as usual it makes no guarantees about their absolute value, only their changes over time. [https://github.com/spulec/freezegun]

from datetime import datetime
from freezegun import freeze_time
@freeze_time(datetime(2022,2,22,12,34,56,789000))
def test_time(self):
    actual = hey_time()
    self.assertEqual("hey: it's 2022-02-22 12:34:56.789000", actual)
def hey_time():
    return 'hey: it's ' + datetime.utcnow()

Case II: How to mock a function to return different responses based on the incoming parameters?

Imagine Trinity and and Neo walking to a cafe, Trinity asks the barista for 2 cup of coffee, one with sugar and the other without sugar.

If you are coming from a Java background, you probably know that you can control the method behavior using when in mocking. Fortunately, python offers similar solution, but with different syntax. So, let’s assume that you have a function that calls S3 bucket and returns a file, the function receives the file name parameter and fetch it from S3 bucket, how can you return different file based on the input parameter? By assigning side_effect property in the mocked function to MagicMock, you will be able to take control over the outcome of the function based on the input parameter. In the code below I created the function handler_side_effect that will examine the first incoming argument and return different result based on that. Now imagine what complicated scenarios you can come up with based on args array.

Trinity doesn’t care what’s going on behind the scene, she cares about the final output.

When we mock the cafe, we treat it as a box, we give it an input and we expect an output.

test.py 
REGION = 'jordan-east-1'

@mock.patch('src.s3_util.fetch_file_from_s3')
def test_handler(self, mock_fetch_file_from_s3):
    mock_fetch_file_from_s3.side_effect = MagicMock(side_effect=handler_side_effect)
def handler_side_effect(*args, **kwargs):
    if args[0] == 'Regular':
        return {"Ingredients ": "Humus, Tomato, Falafel, and Bread"}
    elif args[0] == 'Super':
         return {"Ingredients ": "Humus, Fries, Tahina, Eggplant, Tomato, Falafel, and Bread"}
src/s3_util.py
def fetch_file_from_s3(file_name: str):
    ...
    return file

Case III: How to mock a function to return one value all the time?

If you have a simple where your function needs to return the same value despite the input parameter, then you can use return_value attribute in the mocked function.

@mock.patch('src.s3_util.fetch_file_from_s3')
def test_handler(self, mock_fetch_file_from_s3):
    mock_fetch_file_from_s3.return_value = 'I'm cool as Avocado'

Case IV: How to verify that mocked function was called with specific arguments?

In some cases it’s not enough to verify the outcome of the function you are mocking, you also want to verify whether the mocked function was called with specific arguments? how many times it was called? etc …
It’s important to mention that order of mocked functions is confusing a little bit in python, in the example below, you will notice that mock_humus comes first in the parameters, so keep that in mind to avoid banging your head against the wall :)

@mock.patch('src.eggplant.add')
@mock.patch('src.humus.add')
def test_falafel_factory(self, mock_humus, mock_eggplant):
    mock_humus.assert_called_with('humus')
    mock_humus.assert_called_with('lentils')
    self.assertEqual(mock_humus.call_count, 2)
    self.assertEqual(mock_eggplant.call_count, 0)
src/falafel_sandwitch_factory.py
from humus  import add
def make_falafel_sandwitch(self):
     add('humus')
     add('lentils')

src/humus.py
def add(self, ingredient):
     ...

src/eggplant.py
def add(self, ingredient):
     ...

Case V: How to test your function when exception is raised?

In case your function has to handle an exception, you want to be able to test your function behavior, so let’s say that you have a function that fetches information from s3 bucket and do some analysis on this information, and you are experiencing exception from s3, something like file not found, or region is down. mocking gives you the power to control the behavior of the fetch function, you can raise an exception when the mocked function is called, catch the exception in your code and verify that mocked function was called with certain parameters.

@mock.patch('src.data_factory.read_from_s3')
def test_data_factory(self, mock_read_from_s3):
    mock_read_file.side_effect =    MagicMock(side_effect=Exception('Trinity is mad at you')) 
    try:
       fetch_data_and_analyse()
    except Exception as e: 
       self.assertEqual('Trinity is mad at you', format(error))
    mock_read_from_s3.assert_called_with('cool_bucket')
    self.assertEqual(mock_read_from_s3.call_count, 1)
src/data_factory.py
from s3_util  import read_from_s3
def fetch_data_and_analyse(self):
   try:
     data = read_from_s3('cool_bucket')
   except Exception as e:
     print(e)
   analyse(data)
src/s3_util.py
def read_from_s3(self, bucket_name):
     ...

This is it for now, if you have any case that’s not mentioned in the cases above, please leave me a comment and I will try to get back to you if I’ve the answer.