Python: Mock 常用功能

前言

unittest.mock is a library for testing in Python. It allows you to replace parts of your system under test with mock objects and make assertions about how they have been used.

Python 3.3 新增了一个可以用来在单元测试的时候进行 mock 操作的 unittest.mock 模块。 对于 Python 3.3 之前的版本也可以通过 pip install mock 安装拥有相同功能的移植版模块 mock 。本文记录一些常用的 unittest.mock 模块的用法,理论上同样也适用于 mock 这个第三方模块(注:本文例子是在 Python 3.6 下测试的)。

常见用法

定义返回值

>>> from unittest.mock import MagicMock
>>> mock = MagicMock(return_value=3)
>>> # or
>>> # mock = MagicMock()
>>> # mock.return_value = 3
>>> mock()
3

定义变化的返回值

>>> mock = MagicMock(side_effect=[1, 2, 3])
# or
# >>> mock = MagicMock()
# >>> mock.side_effect = [1, 2, 3]
>>> mock()
1
>>> mock()
2
>>> mock()
3
>>> mock()
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/xxx/lib/python3.6/unittest/mock.py", line 939, in __call__
    return _mock_self._mock_call(*args, **kwargs)
  File "/xxx/lib/python3.6/unittest/mock.py", line 998, in _mock_call
    result = next(effect)
StopIteration


>>> def side_effect(arg=1):
...     return arg
...
>>> m = MagicMock(side_effect=side_effect)
>>> m()
1
>>> m(1)
1
>>> m(2)
2

定义抛出异常

>>> m.side_effect = KeyError
>>> m()
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/xxx/python3.6/unittest/mock.py", line 939, in __call__
    return _mock_self._mock_call(*args, **kwargs)
  File "/xxx/python3.6/unittest/mock.py", line 995, in _mock_call
    raise effect
KeyError
>>>

mock 变量/属性

example.py:

foo = {}


def bar():
    return foo


class FooBar:
    def __init__(self):
        self.msg = 'test'

    def hello(self):
        return 'hello'


def foobar():
    return FooBar().hello()


fb = FooBar()


def hello():
    return fb.msg
>>> from unittest.mock import patch
>>> m = MagicMock()
>>> m.test = MagicMock(return_value=233)
>>> m()
<MagicMock name='mock()' id='4372854824'>
>>> m.test
<MagicMock name='mock.test' id='4372854768'>
>>> m.test()
233
>>>
>>> import example
>>> example.foo
{}
>>> example.hello()
'test'
>>> with patch.object(example, 'foo', {'lalala': 233}):
...     example.foo
...
{'lalala': 233}
>>> example.foo
{}
>>> with patch.object(example.fb, 'msg', 666):
...     example.hello()
...
666
>>> example.hello()
'test'

mock dict

>>> foo = {'a': 233}
>>>
>>> foo['a']
233
>>>
>>> with patch.dict(foo, {'a': 666, 'b': 222}):
...     print(foo['a'])
...     print(foo['b'])
...
666
222
>>> foo['a']
233
>>> 'b' in foo
False
>>> with patch.dict(foo, a=666, b=222):
...     print(foo['a'])
...     print(foo['b'])
...
666
222
>>> foo['a']
233
>>>

mock 函数

>>> import example
>>> from unittest.mock import MagicMock, patch
>>> example.bar()
{}
>>> with patch('example.bar', MagicMock(return_value=233)):
...     example.bar()
...
233
>>>
>>> with patch.object(example, 'bar', MagicMock(return_value=233)):
...     example.bar()
...     example.bar(233)
...
233
233
>>>
>>> with patch.object(example, 'bar', autospec=True) as bar:
...     bar.return_value = 666
...     example.bar()
...
666
>>> with patch.object(example, 'bar', autospec=True):
...     example.bar(233)
...
Traceback (most recent call last):
  File "<stdin>", line 2, in <module>
  File "<string>", line 2, in bar
  File "/xxx/lib/python3.6/unittest/mock.py", line 171, in checksig
    sig.bind(*args, **kwargs)
  File "/xxx/lib/python3.6/inspect.py", line 2934, in bind
    return args[0]._bind(args[1:], kwargs)
  File "/xxx/lib/python3.6/inspect.py", line 2855, in _bind
    raise TypeError('too many positional arguments') from None
TypeError: too many positional arguments

autospec=True 可以保证 mock 后的对象拥有跟原来一样的函数签名。

mock 方法

>>> example.foobar()
'hello'
>>> with patch('example.FooBar') as mocked_cls:
...     instance = MagicMock()
...     instance.hello = MagicMock(return_value='mocked_hello_method')
...     mocked_cls.return_value = instance
...     example.foobar()
...
'mocked_hello_method'
>>>

mock @property 装饰的方法

可以使用 PropertyMock mock @property 装饰的方法以及描述符。

>>> class Foo:
...     def __init__(self):
...             self._bar = 233
...     @property
...     def bar(self):
...             return self._bar
...     @bar.setter
...     def bar(self, value):
...             self._bar = value
...
>>> f = Foo()
>>> f.bar
233
>>> f.bar = 666
>>> f.bar
666
>>> with patch('__main__.Foo.bar', new_callable=PropertyMock) as mock_bar:
...     mock_bar.return_value = 'mocked_bar'
...     foo = Foo()
...     print(foo.bar)
...     foo.bar = 266
...     print(foo.bar)
...
mocked_bar
mocked_bar
>>> mock_bar
<PropertyMock name='bar' id='4373101536'>
>>> mock_bar.mock_calls
[call(), call(266), call()]
>>>

mock 内置函数

test.py:

def to_str(obj):
    return str(obj)
>>> test.to_str('233')
'233'
>>> with patch.object(__builtins__, 'str', int):
...     test.to_str('233')
...     test.to_str('foobar')
...
233
Traceback (most recent call last):
  File "<stdin>", line 3, in <module>
  File "/xxx/test.py", line 2, in to_str
    return str(obj)
ValueError: invalid literal for int() with base 10: 'foobar'
>>>

mock open 函数

read:

>>> from unittest.mock import mock_open
>>> m = mock_open(read_data='mocked data')
>>> with patch('__main__.open', m):
...     with open('test') as fp:
...         print(fp.read())
...
mocked data
>>> m
<MagicMock name='open' spec='builtin_function_or_method' id='4369854192'>
>>> m.mock_calls
[call('test'),
 call().__enter__(),
 call().read(),
 call().__exit__(None, None, None)]
>>>

write:

>>> m2 = mock_open()
>>> with patch('__main__.open', m2):
...     with open('bar', 'w') as fp:
...             fp.write('write data')
...
>>> m2.mock_calls
[call('bar', 'w'),
 call().__enter__(),
 call().write('write data'),
 call().__exit__(None, None, None)]
>>> m2().write.mock_calls
[call('write data')]
>>>

mock isinstance

可以用 __class__ 属性来通过 mock 对象的 isinstance 检查:

>>> m = MagicMock()
>>> m.__class__ = int
>>> isinstance(m, int)
True
>>> m.__class__ = dict
>>> isinstance(m, dict)
True
>>>

mock 真值判断

可以通过 mock __bool__ 方法(Python 3) 或 __nonzero__ 方法(Python 2) 来 mock 对对象的真值判断:

>>> m = MagicMock()
>>> bool(m)
True
>>> m = MagicMock(__bool__=MagicMock(return_value=False))
>>> bool(m)
False
>>> if not m: print('false')
...
false
>>>

mock 特殊属性

所谓的特殊属性就是 Mock/MagicMock 对象自带的属性(自带一些可选参数 [1] ), 比如 name , 如果我们要 mock 的对象也有一个 name 属性的话, 通过 MagicMock(name='foo') 是无法 mock name 这个属性的, 可以用下面两种方法实现 mock 特殊属性:

覆盖 mock 对象的属性:

>>> m = MagicMock(name='foo')
>>> m.name
<MagicMock name='foo.name' id='4372898816'>
>>>
>>> m = MagicMock()
>>> m.name = 'foo'
>>> m.name
'foo'
>>>

使用 configure_mock:

>>> m = MagicMock()
>>> m.name
<MagicMock name='mock.name' id='4373042848'>
>>> m.configure_mock(name='foo')
>>> m.name
'foo'
>>>

自动撤销 mock

patch, patch.object 都支持通过 with 语句或者装饰器可以实现 只在指定函数/方法/块中应用 mock ,函数/方法/块结束后自动撤销 mock

>>> str('abc')
'abc'
>>> with patch.object(__builtins__, 'str', int):
...     str('abc')
...
Traceback (most recent call last):
  File "<stdin>", line 2, in <module>
ValueError: invalid literal for int() with base 10: 'abc'
>>> str('abc')
'abc'
>>>
>>> @patch.object(__builtins__, 'str', int)
... def test():
...     str('abc')
...
>>> str('abc')
'abc'
>>> test()
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/xxx/lib/python3.6/unittest/mock.py", line 1179, in patched
    return func(*args, **keywargs)
  File "<stdin>", line 3, in test
ValueError: invalid literal for int() with base 10: 'abc'
>>>

常用检查方法

mock 的对象拥有一些可以用于单元测试的检查方法,可以用来测试 mock 对象的调用情况。

检查调用次数

待检查的 mock 对象:

>>> m = MagicMock()
>>> m(1)
<MagicMock name='mock()' id='4372904760'>
>>> m(2)
<MagicMock name='mock()' id='4372904760'>
>>> m(3)
<MagicMock name='mock()' id='4372904760'>
>>>

.called: 是否被调用过

>>> m.called
True
>>>

.call_count: 获取调用次数

>>> m.call_count
3

.assert_called(): 检查是否被调用过,如果没有被调用过,则会抛出 AssertionError 异常

>>> m.assert_called()
>>>

.assert_called_once(): 确保调用过一次,如果没调用或多于一次,则抛出 AssertionError 异常

>>> m.assert_called_once()
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/xxx/lib/python3.6/unittest/mock.py", line 795, in assert_called_once
    raise AssertionError(msg)
AssertionError: Expected 'mock' to have been called once. Called 3 times.
>>>

.assert_not_called(): 确保没被调用过,否则抛出 AssertionError 异常

>>> m.assert_not_called()
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/xxx/lib/python3.6/unittest/mock.py", line 777, in assert_not_called
    raise AssertionError(msg)
AssertionError: Expected 'mock' to not have been called. Called 3 times.

检查调用时使用的参数

待检查的 mock 对象:

>>> m = MagicMock()
>>> m(1, 2, foo='bar')
<MagicMock name='mock()' id='4372980792'>

.call_args: 最后一次调用时使用的参数,未调用则返回 None

>>> m.call_args
call(1, 2, foo='bar')
>>>

.assert_called_once_with(*args, **kwargs): 确保只调用过一次,并且使用特定参数调用

>>> m.assert_called_once_with(1, 2, foo='bar')
>>>
>>> m(2)
<MagicMock name='mock()' id='4372980792'>
>>>
>>> m.assert_called_once_with(1, 2, foo='bar')
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/xxx/lib/python3.6/unittest/mock.py", line 824, in assert_called_once_with
    raise AssertionError(msg)
AssertionError: Expected 'mock' to be called once. Called 2 times.
>>>

.assert_any_call(*args, **kwargs): 检查某次用特定参数进行过调用

>>> m.assert_any_call(1, 2, foo='bar')
>>>

.assert_called_with(*args, **kwargs): 检查最后一次调用时使用的参数

>>> m.assert_called_with(2)
>>>
>>> m.assert_called_with(1, 2, foo='bar')
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/xxx/lib/python3.6/unittest/mock.py", line 814, in assert_called_with
    raise AssertionError(_error_message()) from cause
AssertionError: Expected call: mock(1, 2, foo='bar')
Actual call: mock(2)
>>>

.call_args_list: 所有调用时使用的参数列表

>>> m.call_args_list
[call(1, 2, foo='bar'), call(2)]
>>> m(3)
<MagicMock name='mock()' id='4372980792'>
>>> m.call_args_list
[call(1, 2, foo='bar'), call(2), call(3)]

.assert_has_calls(calls, any_order=False): 检查某几次调用时使用的参数, 如果 any_orderFalse 时必须是挨着的调用顺序,可以是中间的几次调用, 为 Truecalls 中的记录可以是无序的

>>> from unittest.mock import call
>>> m.call_args_list
[call(1, 2, foo='bar'), call(2), call(3)]
>>>
>>> m.assert_has_calls([call(2), call(3)])
>>>
>>> m.assert_has_calls([call(3), call(2)])
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/xxx/lib/python3.6/unittest/mock.py", line 846, in assert_has_calls
    ) from cause
AssertionError: Calls not found.
Expected: [call(3), call(2)]
Actual: [call(1, 2, foo='bar'), call(2), call(3)]
>>> m.assert_has_calls([call(3), call(2)], any_order=True)
>>>

.method_calls: mock 对象的方法调用记录

>>> m.test_method(2, 3, 3)
<MagicMock name='mock.test_method()' id='4372935456'>
>>> m.method_calls
[call.test_method(2, 3, 3)]

.mock_calls: 记录 mock 对象的所有调用,包含方法、magic method 以及返回值 mock

>>> m.mock_calls
[call(1, 2, foo='bar'), call(2), call(3), call.test_method(2, 3, 3)]
>>>
>>> m.call_args_list
[call(1, 2, foo='bar'), call(2), call(3)]

手动重置 mock 调用记录

可以使用 .reset_mock() 重置 mock 对象记录的调用记录:

>>> m.mock_calls
[call(1, 2, foo='bar'), call(2), call(3), call.test_method(2, 3, 3)]
>>>
>>> m.call_args_list
[call(1, 2, foo='bar'), call(2), call(3)]
>>>
>>> m.reset_mock()
>>>
>>> m.call_args_list
[]
>>> m.mock_calls
[]
>>>

总结

本文记录的是平时常用的一些 mock 相关的功能,更多内容详见 官方文档


Comments