Use mock in Python unittest

To test an imported package or class, we need to mock it in our test. This post explains how to use some basic functions of mock in Python unit test.

If we have 2 modules, m1 and m2 as below:

# function in m1
def func_m1():
    print("func_m1() in m1.py")

from m1 import func_m1

# function in m2
def func_m2():
    func_m1()
    print("func_m2() in m2.py")

To test m2:

from unittest import mock
from m2 import func_m2

def test_func():
    with mock.patch('m2.func_m1') as mock_func_m1:
        func_m2()
    mock_func_m1.assert_called_once()


if __name__ == "__main__":
    test_func()

When patching func_m1 from m1, we are using m2.func_m1 because func_m1 has been imported into m2.

The result of running code above is:

func_m2() in m2.py

There’s no output from m1 because func_m1 is replaced by mock_func_m1.

The with statement can also be written as decorator:

@patch("m2.func_m1")
def test_func(mock_func_m1):
    func_m2()
    mock_func_m1.assert_called_once()

Mock object has several functions to indicate the number of being called and the variables used in calling. Check python unittest mock document for more information.

Change the action or attribute of a mock

A mock can be assigned with an attribute:

a = Mock()
a.attr_1 = "hi" 
print(a.attr_1)  # hi

A mock can be assigned with an action:

a = Mock()
a.action_1.return_value = "hello"
a.action_1(). # hello

The action of a mock can have side effects.

a = Mock()
# return 3 numbers and throw InterruptedError at the end
a.action.side_effect = [1, 2, 3, InterruptedError]  
try:
    while True:
        print(a.action())
except InterruptedError:
	pass

# result: 1 2 3

Using side_effect can define the action that a function can have every time when it’s called. It’s useful when we test a infinite loop. As code above, we can throw InterruptedError to break the infinite loop.

Mock attribute of a class

If we have a class in m1:

class Alpha():
    a = 1
    def func_alpha(self):
        print(self.a)

To mock attribute a of Alpha:

mock_a = mock.PropertyMock()

mock_a.side_effect = [1, 2, 3]
obj_a = Alpha()
type(obj_a).a = mock_a
obj_a.func_alpha()  # 1
obj_a.func_alpha()  # 2
obj_a.func_alpha()  # 3

Check the calls of mock

We can check if a function of mock is called:

a = Mock()
a.action.return_value = 1
a.action("hi")
a.action.assert_called_once()  # True
assert a.action.call_count == 1
a.action.assert_any_call("hi")
a.actoin.assert_has_calls([call("hi")])

Reset a mock

To reset the calling status of a mock object:

a = Mock()
a.action.side_effect = [1, 2]
a.action()
a.action.assert_called_once()
a.action.reset_mock()
a.action.assert_called_once()  # Assertion error here