python: unit test with mock functions from different modules

I recently started learning python 3 and unit test with pytest and unittest.

As I struggle to figure out how to mock in several situations, I am taking note here so that anyone has same issue maybe find this useful someday.

Structures and code

Before writing tests, this is my folder and files structures.

src/
    ├── my.py  
    ├── my_modules/
    │   ├── __init__.py
    │   └── util.py
    └── tests/
        ├── __init__.py
        ├── test_my.py
        └── test_unit.py

my.py

from my_modules.util import util

def main():
    return util('my input')

if __name__ == '__main__':
    main()

util.py

from datetime import datetime

def util(input: str) -> str:
    input = add_time(input)
    return f"util: {input}"

def add_time(input: str) -> str:
    return f"{datetime.now()}: {input}"

Add a unit test for util function

For util method, I need to mock add_time that is from the same file. I found several ways to achieve this, and this is one of them.

test_unit.py

from unittest.mock import patch, Mock
from my_modules.util import util

@patch('my_modules.util.add_time', Mock(return_value='dummy'))
def test_util():
    expected = 'util: dummy'
    input = 'input'
    input = util(input)
    assert expected == input

I used unittest.mock.patch function as a decorator and specify the desired return value as part of the Mock object. The namespace for add_time is the same as the util function as they are in the same module.

Add a unit test for add_time function

To unit test add_time function, I need to mock datetime.now() function. I again use the unittest.mock.patch. This time, I need to create the Mock with a bit more code as I need to mock a function, rather than the simple return_value nor an attribute.

from datetime import datetime
from unittest.mock import patch, Mock
from my_modules.util import add_time

@patch('my_modules.util.datetime', Mock(**{"now.return_value": datetime(2023, 1, 1)}))
def test_add_time():
    expected = f'{datetime(2023, 1, 1)}: input'
    input = 'input'
    input = add_time(input)
    assert expected == input

I can pass a dictionary that contains attributes and methods information to the Mock object. As I mock now function, I use "now.return_value":<some date>.

If it's an attribute, I can simply pass it to the Mock like Mock(attribute_name=<value>) or as part of the dictionary like {"attribute":<value>}

The module name is a bit interesting as I need to specify my_modules.util.datetime. The reason is that as soon as the datetime is imported to the util.py, it becomes part of the same module, that was quite confusing to me.

I can show another sample of this in the next test.

Add a unit test for main function

To test the main method in the my.py, I just need to mock util method. Let's do it.

test_my.py

from my import main
from unittest.mock import patch, Mock

@patch("my.util", Mock(return_value="dummy"))
def test_main():
    result = main()
    assert result =='dummy'

Even though the util function comes from the my_modules module, in the test time, it becomes my.util namespace as I previously explained.

I specify the Mock object in the decorator, but I cannot if the util function is actually called with the expected parameters. So, let's accept the mock inside the test function.

@patch("my.util")
def test_main_util_called_with_expected_parameter(util_mock):
    util_mock.return_value = 'dummy'
    result = main()
    assert result =='dummy'
    util_mock.assert_any_call('my input')

This time, I use the decorator without passing the Mock object. Instead, I received it in the argument as util_mock, then I specify the return_value.

As I have access to the mock, I can then assert it the method is called with the expected arguments.

Lastly, I also learnt that I can use with statement as well to achieve the same.

def test_main_util_called_with_expected_parameter_with():
    with patch("my.util") as util_mock:
        util_mock.return_value = 'dummy'
        result = main()
    assert result =='dummy'
    util_mock.assert_any_call('my input')

By doing this, I don't need the decorator.

All the tests run successfully!

Summary

I actually don't know yet which is the best way to mock functions. I believe each one has pros and cons, so if any of you have any suggestions or opinions, please let me know in the comment! Thanks.