Python unittest.mock 的基本使用

平时随意用 Python 写的用了就丢的小工具当然没必要写什么单元测试,如果是一个要反复打磨的工具,项目,单元测试就必须重视起来了。因为简单思路写出来的代码将要应对各种未知情形,加之想要大肆的重构,有了足够的测试用例才能安心。

任何语言白盒的单元测试首先面对的就是 Mock, 不光是应对有副作用的操作,最好是隔离所有不在当前所见到的代码的任何调用,其实像 print 这种有副作用的代码一般是能容忍的,可能也是为何很多测试框架默认会把控制台的输出关掉。

在 Python 中,一般基本的东西自己都有,Mock 就直接用 unittest.mock,可应用于所有的单元测试框架中,如 unittest, pytest。另有一块对 uniitest.mock 的简单封装,专为 pytest 提供方便的 pytest-mock。今天先了解 unittest.mock 的使用,它的官方文档是 unittest.mock -- mock object library

unittest.mock 常见的是类 Mock 和 MagicMock, 还有方法(也用作装饰器) patch()。MagicMock(MagicMock(MagicMixin, Mock):) 是增强型的 Mock, 如果愿意的话,可以用 MagicMock g一撸到底。unittest.mock 的功能和其他语言的 Mock 框架一样,可用来 Mock 方法调用和断言属性,调用次数,参数。Python 的 Mock 可比其他静态语言强大的多,能 Mock 的对象几乎不受限,什么类方法,静态方法,还有 Python 自带库的方法,而且使用起来特别灵活。这大概是得益于 Python 的元编程或可随时动态添加方法和属性的特点。
 1>>> from unittest.mock import Mock
 2>>> m1 = Mock()
 3>>> dir(m1)
 4['assert_any_call', 'assert_called', 'assert_called_once', 'assert_called_once_with', 'assert_called_with', 'assert_has_calls', 'assert_not_called', 'attach_mock', 'call_args', 'call_args_list', 'call_count', 'called', 'configure_mock', 'method_calls', 'mock_add_spec', 'mock_calls', 'reset_mock', 'return_value', 'side_effect']
 5>>> m1.method1.return_value = "hello"
 6>>> type(m1.method1)
 7<class 'unittest.mock.Mock'>
 8>>> v1 = m1.method1()
 9>>> v1 'hello'
10>>>
11>>> m1.obj1.obj2.method1.return_value = "world"
12>>> type(m1.obj1)
13<class 'unittest.mock.Mock'>
14>>> type(m1.obj1.obj2.method1)<br><class 'unittest.mock.Mock'>
15>>> m1.method1.assert_called_once()
16>>> m1.method1.assert_not_called()
17Traceback (most recent call last):
18  File "<stdin>", line 1, in <module>
19  File "/opt/homebrew/Cellar/python@3.12/3.12.11/Frameworks/Python.framework/Versions/3.12/lib/python3.12/unittest/mock.py", line 910, in assert_not_called
20    raise AssertionError(msg)
21AssertionError: Expected 'method1' to not have been called. Called 1 times.
22Calls: [call()].
23>>> m1.attr1=100

通过 Mock 或 MagicMock 可以任意(多层次)的添加属性或方法,只要没用 return_value, side_effect, 直接赋值等反回的仍然是一个 Mock。

下面进入到各种实际应用场景

Mock 外部方法调用

下面的代码,比如 repo.get_data() 需要从数据库中查询数据
1def foo(repo):
2   data = repo.get_data()
3   return f"hello {data}"

测试时需要 Mock 住 repo.get_data() 调用,这时我们的测试用例就可以写成
1def test_foo():
2    mock_repo = Mock()  # 明确 repo 的类型可用 Mock(spec=Repo)
3    mock_repo.get_data.return_value = "mocked data"
4
5    result = foo(mock_repo)
6
7    assert result == "hello mocked data"
8    mock_repo.get_data.assert_called_once()

Mock 有副作用外部操作(用到 patch)

比如一个函数进行了文件删除操作,我们当然不想在单元测试中有这样的破坏性操作,这时候要用到 patch。待测试方法如下
1def foo():
2   folder = os.getenv("XYZ")
3   shutil.rmtree(folder, ignore_errors=True)
4   return f"deleted {folder}"

单元测试代码
 1def test_foo():
 2    mock_getenv = patch("os.getenv").start()
 3    mock_getenv.return_value = "world"
 4
 5    mock_rmtree= patch("shutil.rmtree").start()
 6
 7    result = foo()
 8
 9    assert result == "deleted world"
10    mock_getenv.assert_called_once_with("XYZ")
11    mock_rmtree.assert_called_once_with("world", ignore_errors=True)

以上代码可以安全的执行,因为 shutil.rmtree 不会进行实际的操作。如果会删除文件的话基本上测试是不成功的,因为文件系统中大约是没有 "world" 这个目录的。

patch("os.getevn").start() 返回的是一个 MagicMock 类型

patch 用了 `@contextlib.contextmanager` 注解,所以可以写成 with context manager 的形式,上
1def test_foo():
2    with patch("os.getenv") as mock_getenv, patch("shutil.rmtree") as mock_rmtree:
3        mock_getenv.return_value = "world"
4
5        result = foo()
6
7        assert result == "deleted world"
8        ...

这时候就不需要 start() 了

patch 还是一个装饰器,所以又有另一种形式,通过装饰器注入函数参数
1@patch("shutil.rmtree")
2@patch("os.getenv")
3def test_foo(mock_getenv, mock_rmtree):
4        mock_getenv.return_value = "world"
5
6        result = foo()
7
8        assert result == "deleted world"

注意 @patch 时方法参数的顺序是相反的,先写的 @patch 声明为最后的参数。Python 的解释是离方法名最近的装饰器最先读取作为第一个方法参数。如果调个位置,写成
1@patch("shutil.rmtree")
2@patch("os.getenv")
3def test_foo(mock_rmtree, mock_getenv):
4        mock_getenv.return_value = "world"

测试就会失败了,因为上面的 mock_rmtree 实际上是 patch("os.getenv").start() 产生的值。

@patch 还能同时带上 return_value
1@patch("shutil.rmtree")
2@patch("os.getenv", return_value="world")
3def test_foo(mock_getenv, mock_rmtree):

也可 @patch "os", 但必须加上 import os 所在的模块名,例如在 demo.py 中的
1import os
2
3def foo():
4    return os.getenv("XYZ")

测试时就用
1@patch("demo.os")
2def test_foo(mock_getenv):
3    mock_getenv.getenv.return_value = "world"
4
5    assert foo() == "hello world"

Mock 异常

Mock 异常只需要设置 Mock 对象的 side_effect 为一个异常对象即可。比如一个待测试类
 1class Demo:
 2    def foo(self):
 3        try:
 4            return self.bar()
 5        except Exception as ex:
 6            return "wrong: " + str(ex)
 7
 8
 9    def bar(self):
10       # do something, may throw exception
11       return "world"

想要测试在调用 foo() 时,如果其中的 bar() 抛出异常的情况, 我们可以写下面的测试
1def test_demo_foo():
2    demo = Demo()
3    demo.bar = Mock()
4    demo.bar.side_effect = KeyError("no key")
5
6    result = demo.foo()
7
8    assert result == "wrong: 'no key'"

通过这个例子,同时也实现了 Java 中 Mock 对象中部分方法时类似 spy 的功能,即这里的 bar() 被 Mock 了, foo() 还是调用原本的方法。

Mock 特定参数的方法调用

前面用 side_effect 模拟了抛出异常,它同时也可以定制 Mock 方法不同参数时产生不一样的输出,side_effect 可以接受一个 lambda 或函数作为参数。同样是测试
1def foo(env_key):
2    return os.getenv(env_key)

希望 Mock os.getenv() 输入 KEY1,KEY2,和其他时获得不一样的环境变量值,用下面的测试代码
 1from unittest.mock import patch,call
 2
 3def test_foo():
 4    with patch("os.getenv") as mock_getenv:
 5        mock_getenv.side_effect = lambda key: {"KEY1": "value1", "KEY2": "value2"}.get(key, None)
 6        assert foo("KEY1") == "value1"
 7        assert foo("KEY2") == "value2"
 8        assert foo("KEY3") is None
 9
10        mock_getenv.assert_has_calls([call("KEY1"), call("KEY2"), call("KEY3")])

但是如何在别的 key 是调用实际的 os.getenv() 方法呢?我们可以先把原始方法在被 Mock 之前保存下来,然后参考下面的方式
1real_getenv = os.getenv  # real method
2
3def test_foo():
4    with patch("os.getenv") as mock_getenv:
5        mock_getenv.side_effect = lambda key: {"KEY1": "value1", "KEY2": "value2"}.get(key, real_getenv(key)) # call real method for other keys
6        ...

Mock 构造函数

在 Python 中,构造函数可能更多指的是 __new__ 函数,但也可能连带提到 __init__ 初始函数,而我们在用 MyClass() 创建对象时会调用到前两个函数。Mock 构造函数在 Python 中比 Java 等其他语言简单多了,因为 MyClass  本身就可以看作是一个 callable 函数,只要 Mock 它就行了。
1class Cat:
2    ...
3
4class MyClass:
5    def __init__(self):
6        self.cat = Cat()
7
8    def get_cat(self):
9        return self.cat

比如,要 Mock MyClass.__init__ 中的 Cat(), 使用代码
1def test_mock_constructor():
2    with (patch("demo.Cat") as mock_Cat_class):
3        mock_cat = Mock()
4        mock_Cat_class.return_value = mock_cat
5
6        my_class = MyClass()
7
8        assert my_class.get_cat() == mock_cat

unittest.mock 比 Java 的 Mock 框架灵活多了,完全没有构造方法,静态方法,类方法等约束,把它们当作普通方法来 Mock 就行。

Mock dunder 方法也没问题,如上一个测试,加上这样的代码行
1my_class = MyClass()
2mock_cat.__repr__ = Mock(return_value = "mocked Cat __repr__ method")
3
4assert f"{my_class.get_cat()!r} == ""mocked Cat __repr__ method"

所以 Mock __init__ 方法也能尝试一下,从上面的代码来看 Mock dunder 那样的双下划线方法不一定得要 MagicMock, Mock 就能胜任。

patch.object 临时修改某对象的属性或方法

我们无需每次都从最初的函数下手,或对已有 Mock 对象动手脚,针对已有实际对象可用 patch.object(target, attribute,...), 如
1class MyClass:
2    def get_value(self):
3        return "hello"
4
5    def get_value2(self):
6        return "value2"

对已创建的 my_class 实例,我们将只 Mock get_value 方法
1def test_my_class():
2    my_class = MyClass()
3    with patch.object(my_class, "get_value", return_value="world") as mock_get_value:
4        result = my_class.get_value()
5        assert result == "world"
6        mock_get_value.assert_called_once()
7
8        assert my_class.get_value2() == "value2"

即使在 path.object 上下文当中,get_value2() 也是调用实际的方法实现。

把上面第二三行改成如下方式也是一样的效果
1with patch.object(MyClass, "get_value", return_value="world") as mock_get_value:
2    my_class = MyClass()

使用 patch.object(target, attribute,...) 时 target 可以是任何 Python 对象,包括模块本身,类,实例对象,或都是函数, attribute 可以是属性或方法。比如应用到 Mock 前面的 os.getenv 方法
1with patch.object(os, "getenv", return_value="world") as mock_get_value:
2    assert os.getenv("any") == "world"

patch.object() 也可作为装饰器来用,如
1@patch.object(MyClass, "get_value", return_value="world")
2def test_get_value(self, mock_get_data):
3    my_class = MyClass()
4    ...

patch.dict 临时修改 dict 的内容

一看 patch.dict 就应该是 patch.object 专门作用于 dict 快捷方法,看官方的例子
1foo = {}
2@patch.dict(foo, {'newkey': 'newvalue'})
3def test():
4    assert foo == {'newkey': 'newvalue'}
5
6test()
7assert foo == {}

还有更多的 unittest.mock 使用方法,如 patch.multiple(...), mock_open, sentinel, 这些要实际应用时再查询,巩固。

pytest-mock 简介

如果项目中使用 pytest 编写和运行测试用例,建议用 pytest-mock 来 Mock 操作,pytest-mock 是对 unittest.mock 的封装,所以它 100% 兼容后者,学习起来也不难。它主要用借助于 pytest_mock.plugin.MockerFixture 接口来使用,有 Mock 生命周期的自动清理功能,额外提供 spy(), stub(), autospec=True 的功能。

大概看一个比较完整的例子,其中把业务代码与测试代码写在了同一个文件 demo.py 中
 1import unittest
 2import os
 3import requests
 4import sys
 5from pytest_mock.plugin import MockerFixture
 6
 7def get_env_value():
 8    return os.getenv("ENV", "default")
 9
10def fetch_data():
11    env = get_env_value()
12    url = f"https://api.example.com/{env}"
13    response = requests.get(url)
14    return response.text
15
16
17def test_fetch_data_with_mock_and_spy(mocker: MockerFixture):
18    mymodule = sys.modules[__name__]
19
20    mock_get = mocker.patch("demo.requests.get", return_value=mocker.Mock(text="mocked response"))
21    spy_env = mocker.spy(mymodule, "get_env_value")
22
23    result = mymodule.fetch_data()
24
25    assert result == "mocked response"
26    spy_env.assert_called_once()
27    mock_get.assert_called_once_with("https://api.example.com/default")
28
29if __name__ == '__main__':
30    unittest.main()

再应用上 stub() 的话,可以把上面代码的 mock_get = mocker.path(...) 换成以下几行
1    mock_get = mocker.patch("demo.requests.get")
2    fake_response = mocker.stub()
3    fake_response.text = "mocked response"
4    mock_get.return_value = fake_response

会得到一样的效果。

相比于 unittest.mock 的 @patch 前后倒置添加方法参数的方式,pytest-mock 通过单一个参数 mocker: MockerFixture 用起来显得太简洁一些。 本文链接 https://yanbin.blog/python-unittest-mock-usages/, 来自 隔叶黄莺 Yanbin's Blog
[版权声明] 本文采用 署名-非商业性使用-相同方式共享 4.0 国际 (CC BY-NC-SA 4.0) 进行许可。