Python unittest mock 小记

2017/12/19 posted in  Python 黑魔法

自从入坑 TDD 至今,并没有太多的 mock 需求,然而在随着深入到微服务后端组件的开发,依赖的组件越来越多也越来越诡异,便开始真正的正视起 Mock 来。

mock 是 Python 中一个神奇的模块,能够在测试中动态的替换部分逻辑,可以用来解决在跑测试的时候对其他组件依赖的问题。

在 Python 3.3 之后,mock 被加入到 unittest 模块里,不需要单独而在,而在之前,需要单独安装:

pip install mock

对于 mock 模块的学习也非常简单,只需要了解 Mock 类的使用和 Patch 的使用就非常受用了,前者是 mock 的精妙所在,可以创造一个可以替换任意函数、类、对象、方法的 object,并让这个 object 伪装成理想的样子。而 Patch 即是将这个 Mock object 替换的它应该存在的地方。

Mock

要理解 Mock,必须要先理解 Mock 类:

class unittest.mock.Mock(spec=None, side_effect=None, return_value=DEFAULT, wraps=None, name=None, spec_set=None, unsafe=False, **kwargs)

文档描述很详尽,但是常用的参数 side_effect, 和 return_value 值得注意。

side_effect 接受一个函数或异常,当 mock 实例被调用是,便会返回这个函数的返回内容,或者抛出对应的异常:

In [3]: m = Mock(side_effect=lambda: "hello")

In [4]: m()
Out[4]: 'hello'

In [5]: m = Mock(side_effect=Exception("Boooooom"))

In [6]: m()
---------------------------------------------------------------------------
Exception                                 Traceback (most recent call last)
<ipython-input-6-7e5925b669a0> in <module>()
----> 1 m()

....

Exception: Boooooom

当对于返回值比较固定的时候,return_value 可能会更加的实用:

In [7]: m = Mock(return_value="hello")

In [8]: m()
Out[8]: 'hello'

除了在 init 的时候对这些值进行注入,还可以通过赋值的方式:

In [10]: m.func.return_value = "world"

In [11]: m.func()
Out[11]: 'world'

目前只是对 Mock 的简单用法,因为 Mock 的实现很神奇,可以直接调用任何参数、方法,哪怕这些参数、方法并不存在,Mock 实例会直接返回一个新的 Mock 实例,像一棵树一样,非常有趣:

In [13]: m.wat
Out[13]: <Mock name='mock.wat' id='4406809544'>

In [14]: m.wat()
Out[14]: <Mock name='mock.wat()' id='4406812456'>

正是有了这样的特性,我们可以直接进行更深入的修改:

In [17]: m.wat.func.return_value = "watman"

In [18]: m.wat.func()
Out[18]: 'watman'

自此,如何 Mock 类已经非常清晰了,一个 Mock 实例可以作为 callable 的对象存在,而它本身,却又可以包含"无限"层级的新的的 Mock 实例,因此他可以非常简单的伪造出复杂的结构,比如 Response 等。

然而,在 mock 模块中,提供的并不只有一个 Mock 类,以上只是最基本的,还有更自动和常用的 unittest.mock.MagicMock,以及可以替换属性的 unittest.mock.PropertyMock。(除此之外还有两种 NonCallableMock)

MagicMock

MagicMock 是 Mock 的子类,其区别正如其名,即是实现了很多 magic method,可以举一个简单的例子,对于一个实例,如果没有实现 __iter__ 方法时,被转换成 list 会抛出异常:

In [19]: m = Mock()

In [20]: list(m)
---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
<ipython-input-20-150120f9fb24> in <module>()
----> 1 list(m)

TypeError: 'Mock' object is not iterable

如果想正常的转化成列表,那么需要实现这个方法:

In [24]: m.__iter__ = Mock(return_value=iter([]))

In [25]: list(m)
Out[25]: []

MagicMock 类则会自作聪明的 mock 掉所有的 magic method,所以不需要自己去实现,而 __iter__ 默认是空数组:

In [28]: m = MagicMock()

In [29]: list(m)
Out[29]: []

我之所以认为 MagicMock 是自作聪明是因为如果我们需要这些魔法方法,默认的肯定不足以我们使用,我们无法避免自己重写。但因为已经有默认的存在,以至于如果我们手动实现的 Mock 实例存在问题时,往往会被它的默认方法掩盖掉一部分真相。

PropertyMock

PropertyMock 是我又爱又恨的 Mock 类,爱其有用,恨其难写。

在 Mock 类中默认 mock 的是可调用对象,因此,如果是熟悉或者 property 属性的 mock 并不理想:

In [30]: m = MagicMock()

In [31]: m.property_attr
Out[31]: <MagicMock name='mock.property_attr' id='4407739728'>

因此才需要 PropertyMock 的存在来拯救世界:

In [34]: p = PropertyMock(return_value=3)

In [35]: type(m).property_attr = p

In [36]: m.property_attr
Out[36]: 3

这种写法我真是日了狗了,不做详解,荆轲刺秦王。

补充

最后,作为 Mock 类的补充。Mock 除了提供了伪造数据的功能,还保留了大量的自带断言,可以来 ”事后“ 审计被 mock 的方法调用的情况,在单元测试中,简直神器。

Patch

有强大的 Mock 并不够用,只有真正的把伪造的可调用对象替换进去才算完美,而 Patch 便是负责这个工作。

在讲 mock 模块中实现的各种神奇 Patch 之前,依然先提下最基础的 unittest.mock.patch 类:

unittest.mock.patch(target, new=DEFAULT, spec=None, create=False, spec_set=None, autospec=None, new_callable=None, **kwargs)

一目了然,其中 target 是必选项,即使要替换什么,在基础的 Patch 类中,要求为字符串,比如:package.module.ClassName,而 new 即使要替换成的 Mock 类实例,默认是一个 MagicMock 实例,正如前面所说,这个实例实现了大量的默认魔法方法,只能能够让被调用时不抛异常而已,如果需要有数据替换等更精细的操作,需要自己定义 Mock 实例。

Patch 的参数基本上都是作为实例化 Mock 实例时使用的,如果自己手动构造 Mock 实例的话,只需要专注于 Patch 本身的功能就好。可以看一个简单的例子:

In [37]: class Class:
    ...:     def method(self):
    ...:         pass
    ...:

In [38]: with patch('__main__.Class') as MockClass:
    ...:     instance = MockClass.return_value
    ...:     instance.method.return_value = 'foo'
    ...:     assert Class() is instance
    ...:     assert Class().method() == 'foo'

值得一提的是,Patch 可以作为装饰器或者上下文管理器,前者对于默认的 Mock 实例即可满足需求的场景非常实用,而且 Patch 的额外参数也显得非常好用,后者对于构造负责的 Mock 实例非常好用。

Patch 拓展

Patch 类还有三个非常有用的拓展,分别是替换字典用的 patch.dict,多次替换用的 patch.multiple,和替换实例用的 patch.object

前两个非常简单,patch.dict 可以替换字典,或者字典类似实例,参考文档上的例子,一目了然:

>>> foo = {}
>>> with patch.dict(foo, {'newkey': 'newvalue'}):
...     assert foo == {'newkey': 'newvalue'}
...
>>> assert foo == {}

对于 patch.multiple,个人感觉有点鸡肋,目前也没真正的用过,作用正如文档所说:

Perform multiple patches in a single call.

文档的例子:

>>> thing = object()
>>> other = object()
>>> @patch.multiple('__main__', thing=DEFAULT, other=DEFAULT)
... def test_function(thing, other):
...     assert isinstance(thing, MagicMock)
...     assert isinstance(other, MagicMock)
...
>>> test_function()

patch.object 是需要单独拿出来强调的,毕竟和前两者相比,替换一个实例的场景是非常场景的,而且会有非常诡异的替换方式。

其实可以看到,Patch 的参数大同小异,我们只需要关注差异就可以了:

patch.object(target, attribute, new=DEFAULT, spec=None, create=False, spec_set=None, autospec=None, new_callable=None, **kwargs)

首先依然有 target,并且多出了 attribute 参数。对于 Patch 来说,要求 target 是字符串,而这里则要求是要替换的类,反而 attribute 被要求为字符串。除了这点,方法没有其他要注意的了,直接举一个文档的例子:

>>> @patch.object(SomeClass, 'class_method')
... def test(mock_method):
...     SomeClass.class_method(3)
...     mock_method.assert_called_with(3)
...
>>> test()

对于 Python unittest mock 的基础用法算是写完了,可以看到 Mock 和 Patch 都是非常灵性的工具,虽然把实现的奇技淫巧都跟封装了起来,但是其灵活性依然让人惊讶。至于怎么设计 Mock 以应对更复杂的实际场景,应该是另外一种艺术了。