Unittest in python - Mock
Kexin Tang
  2023-07-27 22:11:43 2023-07-27 22:11   2023-07-28 14:48:54    

Intro

Mock is a very useful package for unittest in python. It can replace some classes or functions and change their behaviors, it can also use some built-in methods to help you assert whether pytest calls certain part of your code.

Mock() & MagicMock()

For simplicity, let’s use Mock for example. In most cases, Mock and MagicMock are the same :)

Mock is a class that create all attributes and methods as you access them and store details of how they have been used.

What’s more, you can set anything to a Mock, it will treat them as new Mock (sub Mock).

1
2
3
4
5
6
7
8
9
10
11
12
13
# set an undefined method to a Mock
m = Mock()              # <Mock name="mock">
m.undefined_function()  # <Mock name="mock.undefined_function()">

# use mock as a argument
class Object:
    def func(self, args):
        args.do_something()

o = Object()
m = Mock()
o.func(m)
m.do_something()    # <Mock name="mock.do_something()">

return_value

By setting some methods or functions as Mock, then setting return_value can change original logic: I don’t care about what you write in the function, just return what I want!

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
class Object:
    def __init__(self, x):
        self.x = x
    
    def func(self):
        return self.x

obj = Object(1)

# method 1
obj.func = Mock(return_value = 1024)

# # method 2
# obj.func = Mock()
# obj.func.return_value = 1024

obj.func()  # return 1024 rather than 1

This is always useful in unittest, like:

  1. I don’t want to send a real request via network, just let the requester / dispatcher return what I want;
  2. I don’t want to access a real DB, just tell me what data you have;
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
class MySvc:
    def __init__(self, db):
        self.db = db
        ...

    def myRequest(self, req):
        ...
        results = self.db.fetch(req)
        return results


def test_db(req):
    db = MyDB()
    db.fetch = Mock(return_value = [data1, data2, data3, ...])
    svc = MySvc(db)
    results = svc.myRequest(req)    # [data1, data2, data3, ...]

Sometimes, we will meet some call chains, such as mock.connection.cursor().execute(...).

1
2
3
4
5
6
7
8
9
# mock.call1().call2().call3()
m = Mock()
# get Mock for all calls except the last one
c1 = m.call1.return_value
c2 = c1.call2.return_value
# set Mock for last one call
c2.call3.return_value = "foo"

m.call1().call2().call3()   # "foo"

Basically, we can change the code as following

1
2
3
4
5
6
m = Mock()
c1 = Mock()
m.call1.return_value = c1
c2 = Mock()
c1.call2.return_value = c2
c2.call3.return_value = "foo"

side_effect

side_effect = Exception

1
2
3
4
>> m = Mock()
>> m.exception_side_effect = Mock(side_effect = ValueError)
>> m.exception_side_effect()
ValueError

side_effect = iterable

If we set iterable to side_effect, every time we call it, it will yield one element.

1
2
3
4
5
6
7
8
>> m = Mock()
>> m.iter = Mock(side_effect = [1, 2, 3])
>> m.iter()
1
>> m.iter()
2
>> m.iter()
3

side_effect = callable

1
2
3
4
5
6
7
8
9
def log(*args, **kwargs):
    print(f"args: {args}, kwargs: {kwargs}")

>> m = Mock()
>> m.func = Mock(side_effect = log)
>> m.func()
args: (), kwargs: {}
>> m.func(1, two = 2)
args: (1,), kwargs: {"two": 2}

When we set both return_value and side_effect, the Mock will only use side_effect!!

spec & spec_set

spec can be either a list of string or an existing class / instance. After we set spec, the mock can only have corresponding attributes and methods (just like we use dir to see what attributes and methods does one class support).

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
class Object:
    def __init__(self):
        self.one = 1
        self.two = 2
    def func(self):
        pass

# when we use existing class as spec, the mock hasn't inited
m = Mock(spec = Object)
m.func()    # <Mock name="mock.func()">
m.one       # error, cuz we don't init the Object
m.__init__()
m.one       # 1
m.three = 3 # ok

# when we use existing instance as spec, the mock has inited
o = Object()
m = Mock(spec = o)
m.func()    # <Mock name="mock.func()">
m.one       # 1

# when we use list of string as spec
m = Mock(spec = ["one", "func"])
m.func()    # <Mock name="mock.func()">
m.one       # <Mock name="mock.one">
m.one()     # <Mock name="mock.one()">

The difference between spec and spec_set is, spec can add new stuff while spec_set can only read.

1
2
3
4
5
6
7
m = Mock(spec = ["one"])
m.one
m.two = 2   # ok

mm = Mock(spec_set = ["one"])
m.one
m.two = 2   # error

assertion & call args

Mock supports lots of assertions, such as assert_called, assert_called_once, assert_called_with, etc.

1
2
3
4
m = Mock()
m(1, 2)
m.assert_called()           # True
m.assert_called_with(1, 2)  # True

Mock can also remember what args you used via call_args or call_args_list.

1
2
3
4
5
6
7
m = Mock()
m(1, 2)
m.call_args         # call(1, 2)
m.call_args_list    # [call(1, 2)]
m(3, 4)
m.call_args         # call(3, 4)
m.call_args_list    # [call(1, 2), call(3, 4)]

What’s the difference between these two?

So you can simply think MagicMock = Mock with pre-defined magic methods.

1
2
3
4
>> len(Mock())
Error, Mock doesn't have __len__ method
>> len(MagicMock())
0

So if you want to test or use magic methods in your test, use MagicMock.

If you want to modify the magic methods or just for simplicity purpose, plz use Mock.

  1. 1. Intro
  2. 2. Mock() & MagicMock()
    1. 2.1. return_value
    2. 2.2. side_effect
      1. 2.2.1. side_effect = Exception
      2. 2.2.2. side_effect = iterable
      3. 2.2.3. side_effect = callable
    3. 2.3. spec & spec_set
    4. 2.4. assertion & call args
    5. 2.5. What’s the difference between these two?