在 python 中进行单元测试

单元测试（unit test）就是编写测试来验证某一模块的功能正确性。一般会指定输入，验证输出是否符合预期。

单元测试

进行单元测试，首先要导入 unittest 库。

import unittest

先写一个功能函数，这里以完成加法为例，完成两个数的加法。

def add(a, b):
	return a + b

为了验证加法函数的功能是否正确，首先创建一个 TestAdd 类，继承类 unittest.TestCase，然后在这个类中定义相应的测试函数 test_add()，测试函数要以 test 开头。在函数内部，通常使用 assertEqual()、assertTrue()、assertFalse() 和 assertRaise() 等 assert 断言语句进行验证。

class TestAdd(unittest.TestCase):
	def test_add(self):
		a = 1
		b = 2
		self.assertEqual(a, b, 3)

最后运行测试，在 IPython 和 Jupyter 环境下使用下面的代码。

if __name__ == '__main__'
	unittest.main(argv=['first-arg-ignored], exit=False) 

如果是在命令行下，直接使用 unittest.main() 即可。如果输出 OK，则表示单元测试通过。

上面的例子较为简单，如果比较复杂，可以尝试使用下面的方法，其中心思想就是替换掉被测试函数的一些依赖项。

mock 单元测试

有些情况下函数的依赖很复杂。假设有函数 fun1() 依赖于 fun2() 和 fun3()，而后两者也是很复杂的函数。

from unittest.mock import MagicMock

class M(unittest.TestCase):
	def fun1(self):
		v = self.fun2()
		self.fun3(v)

	def fun2(self):
		...

	def fun3(self):
		...

	def test_fum1(self):
		M = M()
		m.fun2 = MagicMock(return_value=2)
		m.fun3 = MagicMock()
		m.fun1()
		self.assertTrue(m.fun2.called)
		m.fun3.assert_called_with(2)

在这种情况下，可以让 fun2() 替换为一个返回具体的数值，把 fun3() 替换为空函数，这样可以测试 fun1() 调用 fun2()，并用 fun2() 返回值调用 fun3()。

Mock Side Effect

Mock Side Effect 实际上就是 mock 函数，它可以根据不同输入返回不同的值，而不只是 return_value。

def side_effect(arg):
	if arg < 0:
		return 1
	else:
		return 2

mock = MagicMock()
mock.side_effect = side_effect

patch

patch 可以通过装饰器或上下文管理器更方便地使用 mock。

@patch('sort')
def test_sort(self, mock_sort):
	...

这里的 mock_sort 替换 sort，可以像上文一样设置 return_value 和 side_effect。还可以利用 with 语句。

with patch.object(A, '__init__', lambda x: None):
	...

这里用于测试类 A 中地成员函数，它可以避免掉复杂的初始化。