平时随意用 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 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 |
>>> from unittest.mock import Mock >>> m1 = Mock() >>> dir(m1) ['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'] >>> m1.method1.return_value = "hello" >>> type(m1.method1) <class 'unittest.mock.Mock'> >>> v1 = m1.method1() >>> v1 'hello' >>> >>> m1.obj1.obj2.method1.return_value = "world" >>> type(m1.obj1) <class 'unittest.mock.Mock'> >>> type(m1.obj1.obj2.method1)<br><class 'unittest.mock.Mock'> >>> m1.method1.assert_called_once() >>> m1.method1.assert_not_called() Traceback (most recent call last): File "<stdin>", line 1, in <module> 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 raise AssertionError(msg) AssertionError: Expected 'method1' to not have been called. Called 1 times. Calls: [call()]. >>> m1.attr1=100 |
通过 Mock 或 MagicMock 可以任意(多层次)的添加属性或方法，只要没用 return_value, side_effect, 直接赋值等反回的仍然是一个 Mock。
下面进入到各种实际应用场景
Mock 外部方法调用
下面的代码，比如 repo.get_data() 需要从数据库中查询数据
|
1 2 3 |
def foo(repo): data = repo.get_data() return f"hello {data}" |
测试时需要 Mock 住 repo.get_data() 调用，这时我们的测试用例就可以写成
|
1 2 3 4 5 6 7 8 |
def test_foo(): mock_repo = Mock() # 明确 repo 的类型可用 Mock(spec=Repo) mock_repo.get_data.return_value = "mocked data" result = foo(mock_repo) assert result == "hello mocked data" mock_repo.get_data.assert_called_once() |
Mock 有副作用外部操作(用到 patch)
比如一个函数进行了文件删除操作，我们当然不想在单元测试中有这样的破坏性操作，这时候要用到 patch。待测试方法如下
|
1 2 3 4 |
def foo(): folder = os.getenv("XYZ") shutil.rmtree(folder, ignore_errors=True) return f"deleted {folder}" |
单元测试代码
|
1 2 3 4 5 6 7 8 9 10 11 |
def test_foo(): mock_getenv = patch("os.getenv").start() mock_getenv.return_value = "world" mock_rmtree= patch("shutil.rmtree").start() result = foo() assert result == "deleted world" mock_getenv.assert_called_once_with("XYZ") mock_rmtree.assert_called_once_with("world", ignore_errors=True) |
以上代码可以安全的执行，因为 shutil.rmtree 不会进行实际的操作。如果会删除文件的话基本上测试是不成功的，因为文件系统中大约是没有 "world" 这个目录的。
patch("os.getevn").start() 返回的是一个 MagicMock 类型
patch 用了 @contextlib.contextmanager 注解，所以可以写成 with context manager 的形式，上
|
1 2 3 4 5 6 7 8 |
def test_foo(): with patch("os.getenv") as mock_getenv, patch("shutil.rmtree") as mock_rmtree: mock_getenv.return_value = "world" result = foo() assert result == "deleted world" ... |
这时候就不需要 start() 了
patch 还是一个装饰器，所以又有另一种形式，通过装饰器注入函数参数
|
1 2 3 4 5 6 7 8 |
@patch("shutil.rmtree") @patch("os.getenv") def test_foo(mock_getenv, mock_rmtree): mock_getenv.return_value = "world" result = foo() assert result == "deleted world" |
注意 @patch 时方法参数的顺序是相反的，先写的 @patch 声明为最后的参数。Python 的解释是离方法名最近的装饰器最先读取作为第一个方法参数。如果调个位置，写成
|
1 2 3 4 |
@patch("shutil.rmtree") @patch("os.getenv") def test_foo(mock_rmtree, mock_getenv): mock_getenv.return_value = "world" |
测试就会失败了，因为上面的 mock_rmtree 实际上是 patch("os.getenv").start() 产生的值。
@patch 还能同时带上 return_value
|
1 2 3 |
@patch("shutil.rmtree") @patch("os.getenv", return_value="world") def test_foo(mock_getenv, mock_rmtree): |
也可 @patch "os", 但必须加上 import os 所在的模块名，例如在 demo.py 中的
|
1 2 3 4 |
import os def foo(): return os.getenv("XYZ") |
测试时就用
|
1 2 3 4 5 |
@patch("demo.os") def test_foo(mock_getenv): mock_getenv.getenv.return_value = "world" assert foo() == "hello world" |
Mock 异常
Mock 异常只需要设置 Mock 对象的 side_effect 为一个异常对象即可。比如一个待测试类
|
1 2 3 4 5 6 7 8 9 10 11 |
class Demo: def foo(self): try: return self.bar() except Exception as ex: return "wrong: " + str(ex) def bar(self): # do something, may throw exception return "world" |
想要测试在调用 foo() 时，如果其中的 bar() 抛出异常的情况, 我们可以写下面的测试
|
1 2 3 4 5 6 7 8 |
def test_demo_foo(): demo = Demo() demo.bar = Mock() demo.bar.side_effect = KeyError("no key") result = demo.foo() assert result == "wrong: 'no key'" |
通过这个例子，同时也实现了 Java 中 Mock 对象中部分方法时类似 spy 的功能，即这里的 bar() 被 Mock 了, foo() 还是调用原本的方法。
Mock 特定参数的方法调用
前面用 side_effect 模拟了抛出异常，它同时也可以定制 Mock 方法不同参数时产生不一样的输出，side_effect 可以接受一个 lambda 或函数作为参数。同样是测试
|
1 2 |
def foo(env_key): return os.getenv(env_key) |
希望 Mock os.getenv() 输入 KEY1，KEY2，和其他时获得不一样的环境变量值，用下面的测试代码
|
1 2 3 4 5 6 7 8 9 10 |
from unittest.mock import patch,call def test_foo(): with patch("os.getenv") as mock_getenv: mock_getenv.side_effect = lambda key: {"KEY1": "value1", "KEY2": "value2"}.get(key, None) assert foo("KEY1") == "value1" assert foo("KEY2") == "value2" assert foo("KEY3") is None mock_getenv.assert_has_calls([call("KEY1"), call("KEY2"), call("KEY3")]) |
但是如何在别的 key 是调用实际的 os.getenv() 方法呢？我们可以先把原始方法在被 Mock 之前保存下来，然后参考下面的方式
|
1 2 3 4 5 6 |
real_getenv = os.getenv # real method def test_foo(): with patch("os.getenv") as mock_getenv: mock_getenv.side_effect = lambda key: {"KEY1": "value1", "KEY2": "value2"}.get(key, real_getenv(key)) # call real method for other keys ... |
Mock 构造函数
在 Python 中，构造函数可能更多指的是 __new__ 函数，但也可能连带提到 __init__ 初始函数，而我们在用 MyClass() 创建对象时会调用到前两个函数。Mock 构造函数在 Python 中比 Java 等其他语言简单多了，因为 MyClass 本身就可以看作是一个 callable 函数，只要 Mock 它就行了。
|
1 2 3 4 5 6 7 8 9 |
class Cat: ... class MyClass: def __init__(self): self.cat = Cat() def get_cat(self): return self.cat |
比如，要 Mock MyClass.__init__ 中的 Cat(), 使用代码
|
1 2 3 4 5 6 7 8 |
def test_mock_constructor(): with (patch("demo.Cat") as mock_Cat_class): mock_cat = Mock() mock_Cat_class.return_value = mock_cat my_class = MyClass() assert my_class.get_cat() == mock_cat |
unittest.mock 比 Java 的 Mock 框架灵活多了，完全没有构造方法，静态方法，类方法等约束，把它们当作普通方法来 Mock 就行。
Mock dunder 方法也没问题，如上一个测试，加上这样的代码行
|
1 2 3 4 |
my_class = MyClass() mock_cat.__repr__ = Mock(return_value = "mocked Cat __repr__ method") assert f"{my_class.get_cat()!r} == ""mocked Cat __repr__ method" |
所以 Mock __init__ 方法也能尝试一下，从上面的代码来看 Mock dunder 那样的双下划线方法不一定得要 MagicMock, Mock 就能胜任。
patch.object 临时修改某对象的属性或方法
我们无需每次都从最初的函数下手，或对已有 Mock 对象动手脚，针对已有实际对象可用 patch.object(target, attribute,...), 如
|
1 2 3 4 5 6 |
class MyClass: def get_value(self): return "hello" def get_value2(self): return "value2" |
对已创建的 my_class 实例，我们将只 Mock get_value 方法
|
1 2 3 4 5 6 7 8 |
def test_my_class(): my_class = MyClass() with patch.object(my_class, "get_value", return_value="world") as mock_get_value: result = my_class.get_value() assert result == "world" mock_get_value.assert_called_once() assert my_class.get_value2() == "value2" |
即使在 path.object 上下文当中，get_value2() 也是调用实际的方法实现。
把上面第二三行改成如下方式也是一样的效果
|
1 2 |
with patch.object(MyClass, "get_value", return_value="world") as mock_get_value: my_class = MyClass() |
使用 patch.object(target, attribute,...) 时 target 可以是任何 Python 对象，包括模块本身，类，实例对象，或都是函数， attribute 可以是属性或方法。比如应用到 Mock 前面的 os.getenv 方法
|
1 2 |
with patch.object(os, "getenv", return_value="world") as mock_get_value: assert os.getenv("any") == "world" |
patch.object() 也可作为装饰器来用，如
|
1 2 3 4 |
@patch.object(MyClass, "get_value", return_value="world") def test_get_value(self, mock_get_data): my_class = MyClass() ... |
patch.dict 临时修改 dict 的内容
一看 patch.dict 就应该是 patch.object 专门作用于 dict 快捷方法，看官方的例子
|
1 2 3 4 5 6 7 |
foo = {} @patch.dict(foo, {'newkey': 'newvalue'}) def test(): assert foo == {'newkey': 'newvalue'} test() assert 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 中
|
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 27 28 29 30 |
import unittest import os import requests import sys from pytest_mock.plugin import MockerFixture def get_env_value(): return os.getenv("ENV", "default") def fetch_data(): env = get_env_value() url = f"https://api.example.com/{env}" response = requests.get(url) return response.text def test_fetch_data_with_mock_and_spy(mocker: MockerFixture): mymodule = sys.modules[__name__] mock_get = mocker.patch("demo.requests.get", return_value=mocker.Mock(text="mocked response")) spy_env = mocker.spy(mymodule, "get_env_value") result = mymodule.fetch_data() assert result == "mocked response" spy_env.assert_called_once() mock_get.assert_called_once_with("https://api.example.com/default") if __name__ == '__main__': unittest.main() |
再应用上 stub() 的话，可以把上面代码的 mock_get = mocker.path(...) 换成以下几行
|
1 2 3 4 |
mock_get = mocker.patch("demo.requests.get") fake_response = mocker.stub() fake_response.text = "mocked response" mock_get.return_value = fake_response |
会得到一样的效果。
相比于 unittest.mock 的 @patch 前后倒置添加方法参数的方式，pytest-mock 通过单一个参数 mocker: MockerFixture 用起来显得太简洁一些。
本文链接 https://yanbin.blog/python-unittest-mock-usages/, 来自 隔叶黄莺 Yanbin Blog
[版权声明] 
本文采用 署名-非商业性使用-相同方式共享 4.0 国际 (CC BY-NC-SA 4.0) 进行许可。