Python中使用unittest做单元测试时如何优雅地处理input

一、问题描述

假设我们有下面的一个函数需要接受测试：

def add_3(num):
    return 3 + num

上面这个函数就是把给定的参数加上3，然后返回。用unittest来做测试非常简单，代码如下：

import unittest
from operation import add_3


class TestOperation(unittest.TestCase):
    def test_add_3(self):
        # 断言add_3(5)的返回值是8
        self.assertEqual(add_3(5),8)



if __name__ == '__main__':
    unittest.main()

测试这样的函数很简单，很清晰。关键的代码就是self.assertEuqal(add_3(5),8).这行代码的可读性非常强，注释都显得多余。

下面我要修改一下add_3函数，让它不需要接受参数num，而是等待用户输入，根据用户输入的值来返回结果。修改后的代码如下：

def add_3():
    while True:
        num = input('请输入一个整数>>>')
        if num.isdigit():
            break
    return int(num) + 3

改成这样之后，怎么测试呢？add_3不需要参数了，那我们就不给它参数了呗。把之前的测试稍微修改一下，可以变成下面的样子：

 

import unittest
from operation import add_3


class TestOperation(unittest.TestCase):
    def test_add_3(self):
        # 断言add_3()在用户输入5的时候返回值是8
        self.assertEqual(add_3(),8)



if __name__ == '__main__':
    unittest.main()

 如果我们运行这个新的测试，运行到input语句的时候，控制台会停下，等我们输入。为了按照我们的计划完成测试，我们必须输入5。运行效果如下图所示：

测试也顺利地完成了，但是这种方式好吗？很显然，不好。主要的缺陷如下：

1. 不够自动化。这里只测试了一个函数，只需要一次输入。如果同时测试多个函数，每个函数都需要多次输入呢？而且别忘了，输入什么值还是有要求的（比如，这里的计划就是输入5）。这样看来，效率太低。
2. 代码的可读性不高。我们第一个版本的测试语句self.assertEqual(add_3(5),8)做到了完美的自我解释，不需要注释也能看到要测试的数据是什么，期望的结果是什么。但是第二个版本self.assertEuqal(add_3(),8)看起来就像在说这个函数永远返回8一样，必须依赖注释才知道用户输入5的时候这个程序返回8.

幸运的是，unittest已经给我们提供了处理input的好方法。下面我们就来看看如何使用unittest.mock.patch来优雅地实现绕过手动输入的测试。

二、使用unittest.mock.patch来优雅地处理input

patch是什么？patch是一个带参数的装饰器，它可以用来装饰函数（方法），装饰类，还可以当做上下文管理器来使用。我们下面就会用它来装饰方法。

patch有什么用？它会用一个新对象覆盖旧对象，并且在离开了特定的作用域后，取消覆盖，把旧对象还给程序。这样解释有些抽象，结合我们当前的需求来看就会清晰很多。

我们的函数add_3里面用到了input，但其实真正有用的只是input的返回值。因为我们的测试就是基于这个返回值（我们期望的返回值是'5'）。可是这个返回值又依赖于用户的输入，所以就导致程序会停下来，程序员需要输入5，然后测试才能继续进行。如果能绕开“程序暂停”、“程序员手动输入5”这两个步骤，直接得到input的返回值，问题就迎刃而解了。所以，我们就需要用patch来帮我们覆盖。这里，input就是要覆盖的旧对象，而新对象会由patch提供给我们。这个新对象是MagicMock对象，它的属性很多，我们先来看看马上就要使用的return_value。

return_value，顾名思义，就是返回值。我们规定了新对象的返回值是'5'，这就相当于让input的返回值是5.而新对象不需要暂停和输入，所以我们的测试就可以自动进行，不需要手动干预。请看下面的代码实现：

import unittest
from unittest.mock import patch
from operation import add_3


class TestOperation(unittest.TestCase):
    # 传入的参数是说要覆盖的旧对象是内置的input函数
    @patch('builtins.input')
    # 第二个参数就是新对象
    def test_add_3(self,mock_input):
        # 模拟接收用户输入并返回'5'
        mock_input.return_value = '5'
        # 断言用户输入5时，add_3函数会返回8
        self.assertEqual(add_3(),8)



if __name__ == '__main__':
    unittest.main()

还有一点要注意：patch中说明要覆盖的旧对象时，一定要写清楚命名空间。如果这里只写一个input，会报错，因为没有说明input的命名空间。

确定命名空间的原则是：在哪里使用这个对象，哪里就是它的命名空间。所以，假设我们现在所在的文件是b.py,而函数func是a.py中定义的。我们在b.py的头部写了一个from a import func。那么在覆盖func的时候，patch中不能写'a.func'，因为我们不是在a中使用func的，而是在b中使用。这时应该写'__main__.func'（运行时的命名空间已经是__main__了）。如果我们是导入模块a，而不是单独导入func，那么func的命名空间依然是a，覆盖的时候patch中要写‘a.func’。

我们现在要覆盖的是内置函数input，可以写成'builtins.input'，也可以写'__main__.input',前者更清晰。

三、进一步思考，如果要测试的函数中有多个input语句怎么办？

 我们再来仔细看看mock_input.return_value='5'这个语句。其实它的意思可以用下面的代码来表示(具体的实现细节有很大区别，但是可以这样简单地理解)：

class MagicMock:
    def __call__(self):
        return '5'

mock_input = MagicMock()

我们用mock_input覆盖了原来的input,所以程序碰到input语句，就会调用mock_input。而mock_input这个函数不需要参数，并且返回值是固定的'5'。因此看起来就好像用户输入了5一样。

对于只有一个input语句的被测试函数来说，我们只需要一个来自input的返回值，上面的这种实现方式足够了。但如果被测试的函数有两个input语句呢？看看下面的函数：

def add():
    while True:
        first = input('请输入第一个数>>>')
        if first.isdigit():
            break
    while True:
        second = input('请输入第二个数>>>')
        if second.isdigit():
            break
    return int(first) + int(second)

这回是两个数相加了，程序要从用户获取两次输入。如果我们还是用永远返回5的mock_input来覆盖input，就只能僵化地测试5+5=10. 如果我要测试3 + 5 = 8 呢？看来现在这个永远返回一个固定值的mock_input已经不能满足我们的需求了。我们需要一个第一次返回3，第二次返回5的可调用对象。这就要用到MagicMock对象的side_effect属性了。

side_effect属性的值可以是一个函数，可以是一个可迭代对象，也可以是一个异常。为了解决我们的问题，可迭代对象非常合适。如果我们设置mock_input.side_effect = ['3', '5']，那么第一次调用mock_input就会返回'3',第二次调用会返回'5'。这种调用一次就返回一个值，并且每个值还不一样的行为是不是很熟悉？它的背后就是一个迭代器。

请看代码实现：

import unittest
from unittest.mock import patch
from operation import add


class TestOperation(unittest.TestCase):
    @patch('builtins.input')
    def test_add(self,mock_input):
        # 模拟接收用户输入并返回'3','5'
        mock_input.side_effect = ['3','5']
        # 断言用户输入3 5的时候，返回8
        self.assertEqual(add(),8)
    

if __name__ == "__main__":
    unittest.main()

 

如果这篇博文帮到了你，就请给我点个赞吧(#^.^#)

有疑问也欢迎留言~博主可nice啦，在线秒回ヾ(◍°∇°◍)ﾉﾞ