Mock 模块使用说明

功能介绍

好的编码习惯都应该为每一行代码做覆盖测试，但有些时候代码处理的是从网络上获取的内容，或者设备的返回，比如获取交换机路由器的运行结果，或者从网络上获取页面等等。这些动作要么需要联网，要么需要设备，但实际上我们只是想测试代码正确性而已，注重的是对返回的内容的处理而不必非要有实际设备。

mock 模块用于在单元测试中模拟其它代码的结果，比如某个函数需要调用其他函数，这个时候我们可以模拟这个第三方函数的结果来略过实际调用它，不光可以节省时间，也可以避免因为第三方函数出错而影响自己的代码，甚至可以很轻松的模拟难以出现的各种情况。

也正是因为这个模块是如此好用，在 Python2 中还需要单独安装 mock 模块，而 Python3.3 开始这个模块就被放入标准模块了，名叫 unittest.mock

使用思路和实例

在概念上， mock 用于模拟函数的返回，比如你有一个函数调用了另一个函数，而另一个函数的代码本身不是你写的，或者不需要在当前单元测试中测试，你只是希望拿到另一个函数返回的结果，这个时候就可以用 mock 来模拟那个函数来略过各种中间过程而直接得到结果。比如下面这样的代码结构：

                                                                +======================+
                                                           +----|    send_shell_cmd    |
+==========================+    +=====================+    |    +======================+
| test_search_flow_session |----| search_flow_session |----+
+==========================+    +=====================+    |    +======================+
                                                           +----| get_all_flow_session |
                                                                +======================+

上面的 test_search_flow_session 是写在单元测试脚本中的测试案例，用来测试在另一个源代码文件中的 search_flow_session 函数。而 search_flow_session 要调用另 2 个其它文件中的函数 send_shell_cmd 和 get_all_flow_session 来完成功能。恰恰麻烦的是这 2 个函数其中一个需要一台 PC 机来执行 linux 命令，另一个需要一台昂贵的设备来获取设备上的状态和返回，更别说创建拓扑和恢复测试环境的工作，仅仅为了检查 search_flow_session 中的某些代码而付出这样的代价完全不值。

但是应该怎么用 mock 模拟，或者怎么把 mock 注入到你自己的函数中却是一个很伤脑筋的问题，不同的代码风格很容易把你带进坑里，比如要调用的其他函数使用 OOP 方式写的，你会想难道我还得先实例化？或者我的函数是面向对象的，调用的却是面向过程的，怎么办？在我刚刚开始接触 mock 的时候，这些概念和行为真是把我折磨的够呛。写多了之后才慢慢感觉到了下面几个规则：

1. 不用管自己的函数怎么写， mock 只用来模拟别人的模块，不管是面向过程还是面向对象都不用过多考虑，只考虑你的代码中调用了哪些外部函数或者方法，这意味着你要 mock 多少东西
2. 如果调用的外部代码是面向过程的风格，也就是一个一个函数，那么用 mock.patch ；面向对象风格，比如你调用的只是一个类中的某个方法则用 mock.patch.object 。现在看到什么 mock.patch , mock.patch.object 可能你不理解，没事，先放下，到后面会专门说

mock 概念很绕，但是真正用到的接口并不多。也是，模拟函数或者方法行为而已，又能有几种接口呢……大致说来我们能接触到的也就是这么几个：

Mock

mock 是最初，也是最基本的一个函数，它的任务就是模拟某个模块的函数。

patch - 补丁方式模拟

有些函数可能不属于你，你也不在意它的内部实现而只是想调用这个函数然后得到结果而已，这种时候就可以用 patch 方式来模拟。

比如一个模块 linux_tool.py 里面有多个函数，其中 send_shell_cmd 是其他人写的。它具体怎么做我不在乎，只知道它向 Linux PC 发命令然后将命令的结果返回给我。现在我写了一个函数 check_cmd_response 检查返回结果，然后对 check_cmd_response 做单元测试。因为 send_shell_cmd 函数需要一个真实的 PC ，这需要设备且每次返回还可能与预期不符，比如设备无法连接，想检查的东西忘记配置所以取不回来等等，这些都会干扰我自己函数的行为，而且问题和自己函数无关，这种时候就可以用 mock 模拟 send_shell_cmd 函数而且把预期返回写到这个模拟过程中，保证每次都会正确处理。当然有人说可能的确有错误情况啊，这也是你应该要处理的，或者有多种返回啊……没错，所以可以多写几个测试案例把这些情况都模拟一遍嘛。

面向过程代码风格

下面是完整的模拟代码，首先是 linux_tool.py 文件，里面 2 个函数， send_shell_cmd 直接返回一个字符串，注意在现实中这是一个完整函数会连接设备并获取返回的。另一个就是自己写的函数了，中间的代码都去掉，但是整体来说我希望获取未来使用 mock 模拟的函数所返回的内容

#!/usr/bin/env python3
import re

def send_shell_cmd():
    return "Response from send_shell_cmd function"

def check_cmd_response():
    response = send_shell_cmd()
    print("response: {}".format(response))
    return re.search(r"mock_send_shell_cmd", response)

然后是单元测试，注意 patch 的用法，它是一个装饰器，需要把你想模拟的函数写在里面，然后在后面的单元测试案例中为它赋一个具体实例，再用 return_value 来指定模拟的这个函数希望返回的结果就可以了，后面就是正常单元测试代码。

#!/usr/bin/env python3
from unittest import TestCase, mock
import linux_tool

class TestLinuxTool(TestCase):
    def setUp(self):
        pass

    def tearDown(self):
        pass

    @mock.patch("linux_tool.send_shell_cmd")
    def test_check_cmd_response(self, mock_send_shell_cmd):
        mock_send_shell_cmd.return_value = "Response from emulated mock_send_shell_cmd function"

        status = linux_tool.check_cmd_response()
        print("check result: %s" % status)
        self.assertTrue(status)

好了，我们再来梳理一下思路，使用 mock 其实代码方面并没有太多麻烦的，但是厘清思路往往很困难：

1. 实际测试代码和单元测试代码是分开在 2 个文件中的，第一个关卡往往就是怎么把这 2 个文件有机结合起来。这里的关键就是：源代码该怎么写就怎么写，不需要考虑为 mock 留下什么接口之类的东西。

2. 单元测试文件中，首先写单元测试代码，就和正常的一样，最开始的时候只需要 import mock 模块即可。

3. 判断要测试的函数中是否用了其他函数，有可能使用了多个外部函数，那么就判断哪个函数适合 mock ，哪些不需要，一般像浪费时间的，结果不定的，需要其他设备的函数最好都 mock ，其它一些功能函数可用可不用。

4. 确定了哪些外部函数要 mock 就用 patch 语句将它们列出来，每个 patch 是一个函数，而且要确定这些外部函数都在文件头部用 import 语句载入到内存了，因为 mock 模块是通过替换内存中的函数微代码来实现功能的。

5. 如果 patch 多个外部函数，那么调用遵循自下而上的规则，比如：

   @mock.patch("function_C")
   @mock.patch("function_B")
   @mock.patch("function_A")
   def test_check_cmd_response(self, mock_function_A, mock_function_B, mock_function_C):
       mock_function_A.return_value = "Function A return"
       mock_function_B.return_value = "Function B return"
       mock_function_C.return_value = "Function C return"
   
       self.assertTrue(re.search("A", mock_function_A()))
       self.assertTrue(re.search("B", mock_function_B()))
       self.assertTrue(re.search("C", mock_function_C()))

如果函数是在其它文件中实现的，那么 mock 的方式又有不同：

# run_multiple 是在另一个文件 utils.py 中实现的
def run_multiple():
    pass

# 但是在 tool.py 文件中调用了这个模块
from utils import run_multiple

def tool():
    run_multiple()

# test_tool.py 测试的时候就不能 mock 原始实现的路径，而是使用的路径
import unittest2 as unittest
import mock

@mock.patch("tool.run_multiple")
def test_tool(mock_run_multiple):
    mock_run_multiple.return_value = None

上面的关键就是 mock.patch 的路径必须是 "tool.run_multiple" ，这是使用 run_multiple 函数的路径，而不是实现这个函数的路径 "utils.run_multiple"

面向对象代码风格

如果你的代码风格是面向对象的呢？也可以，用 patch.object 就行，来看看例子：

# linux_tool.py
import re

class LinuxTool(object):
    def __init__(self):
        pass

    def send_shell_cmd(self):
        return "Response from send_shell_cmd function"

    def check_cmd_response(self):
        response = self.send_shell_cmd()
        print("response: {}".format(response))
        return re.search(r"mock_send_shell_cmd", response)

再来写单元测试的案例：

from unittest import TestCase, mock
from linux_tool import LinuxTool

class TestLinuxTool(TestCase):
    def setUp(self):
        self.linux_tool = LinuxTool()

    def tearDown(self):
        pass

    @mock.patch.object(LinuxTool, "send_shell_cmd")
    def test_check_cmd_response(self, mock_send_shell_cmd):
        mock_send_shell_cmd.return_value = "Response from emulated mock_send_shell_cmd function"

        status = self.linux_tool.check_cmd_response()
        print("check result: %s" % status)
        self.assertTrue(status)

面向对象的 mock 和面向过程的很相似，唯一就是把 mock.patch 替换成 mock.patch.object ，并且在里面列出类实例和方法名。仔细观察，是类的实例 (不是字符串) 和方法名 (是字符串的方法名而不是方法对象)

side_effect

side_effect 是 mock 中角色比较复杂的方法，它有好几种用法

模拟同一个函数被多次调用

如果要多次调用相同函数并获取返回，比如有一个外部方法叫 linux_tool.send_shell_cmd 用来执行命令并返回命令中间的输出，利用这个函数我又写了一个自己的方法用来建立 vsftpd 服务器，其中多次调用外部方法来创建备份文件，建立配置文件，重启服务，检查服务状态等等。或者某个命令在一个循环中被调用，循环次数也可能是不定的。上面的例子都只是模拟了一次，那么模拟多次怎么办？

答案就是使用 side_effect ，比如下面的例子中在方法 start_ftp_service 中调用了 5 次 send_shell_cmd 方法：

class TestSetupServer(TestCase):
    @mock.patch.object(linux_tool, "send_shell_cmd")
    def test_start_ftp_service_for_default_conf(self, mock_send_shell_cmd):
        mock_send_shell_cmd.side_effect = [
            "cmd1_response",
            "cmd2_response",
            "cmd3_response",
            "cmd4_response",
            "cmd5_response",
        ]

        self.mytool.start_ftp_service()

如果某个命令在循环中被调用，满足判断结果才会跳出循环，那么也要用 side_effect 来模拟循环中的每次结果，一定数清楚具体的循环次数或者精心设计返回，否则执行会出错。

模拟异常

用上面模拟同一个函数多次被调用的实例为例，如果希望主动引发异常，比如 Exception 那么可以这样：

mock_send_shell_cmd.side_effect = Exception("Raise Exception")

所有 raise 语句可以引发的异常都可以用 side_effect 引发

模拟对象中的属性

有些时候要模拟的不是其它类中的方法，而是属性，比如下面这个类里面有一个属性 before ，一个方法 spawnu ，方法的模拟很简单在上面已经有说明，但 before 这个属性呢？这就要用到 mock.PropertyMock 组件了，看下面的例子

class pexpect(object):
    """Fake pexpect class"""
    def __init__(self):
        """INIT"""
        self.before = None

    def spawnu(self):
        """Fake method"""
        pass


class UnitTest(unittest.TestCase):
    @mock.PropertyMock(pexpect, "before")
    @mock.patch.object(pexpect, "spawnu")
    def test_send_cli_cmd(self, mock_spawnu, mock_before):
        pass

MagicMock

mock.MagicMock 是 mock.Mock 的子类，区别就是 MagicMock 预置了其它 MagicMethod ，所谓 MagicMethod 在 Python 中表现为双下划线包围的方法，比如最熟悉的 init 或者 str 之类的。 mock.Mock 默认没有实现这些方法，如果想测试这些方法的行为就得自己写，而 MagicMock 默认预置了这些行为，这样像自增自减，列表的循环，计算符号的重载等 MagicMethod 就在 MagicMock 中内置了，如果不考虑这些那么 MagicMock 和 Mock 行为是一样的

一般情况下模拟都用 MagicMock ，因为这个模拟出来的行为更类似于我们预期

精准模拟第三方函数

自己写的模块大多数时候都需要调用其它函数 (比如大多数模块都会用的 os 或者 sys 模块) ，如何模拟这些第三方函数呢？可以看一个例子：

# 功能模块， 模块名 demo.py
import os

class Demo(object):
    def __init__(self):
        pass

    def delete_file(self, filepath):
        if os.path.isfile(filepath):
            os.remove(filepath)
        return True


# 测试代码，文件名 test_demo.py
from demo import Demo


import mock
import unittest

class TestDemo(unittest.TestCase):
    def setUp(self):
        self.ins = Demo()

    def tearDown(self):
        pass

    @mock.patch("demo.os.path.isfile")
    @mock.patch("demo.os.remove")
    def test_delete_file(self, mock_remove, mock_isfile):
        filepath = "~/tmp/aa"
        mock_isfile.return_value = True
        mock_remove.return_value = True
        self.ins.delete_file(filepath)
        self.assertTrue(mock_remove.called)
        mock_remove.assert_called_with(filepath)
        mock_isfile.assert_called_with(filepath)

上面的例子中 Demo 载入了第三方模块 os ，这个模块很可能在很多模块中都被载入和调用过，如果源码文件特别多的话可能 os 这个模块会到处都是，而测试代码中如果直接模拟 os 模块的话很可能多个 test_ 源文件会互相影响。最好的办法就是对每个源文件的第三方模块精准模拟

在 demo.py 文件中调用了 os.path.isfile 和 os.remove 方法，如何精准模拟呢？上面的例子中用 mock.patch("demo.os...") 的方式就可以做到

内建的其他方法

called

一旦 mock 被创建，比如上面用 patch 模拟的 mock_send_shell_cmd ，或者用 MagicMock 模拟的 mock_func ，都可以用 called() 方法来检查自己究竟有没有被调用，比如：

mock_send_shell_cmd.called
>> True

call_count

返回模拟的函数或方法被调用了几次：

mock_send_shell_cmd.call_count
>> 2

call_args

返回 mock 的东西在调用时传入的具体参数

>>> mock_send_shell_cmd.some_method3(cmd="ls -l", mode="shell")
>>> mock_send_shell_cmd.some_method3.call_args
call(cmd="ls -l", mode="shell")

还有一个叫 call_args_list ，这个用于 mock 的方法被多次调用的情况，会返回一个列表，列表中是每次被调用时的参数

assert_called_with

有时候我们不光想确认自己 mock 的东西有没有被调用，还想确认调用时传入的参数是不是正确的，就可以用 assert_called_with ，比如：

>>> mock_send_shell_cmd.some_method3(a=1, b=4)
>>> mock_send_shell_cmd.some_method3.assert_called_with(a=1, b=4)
>>> mock_send_shell_cmd.some_method3.assert_called_with(a=1, b=5)
Traceback (most recent call last):
...
raise AssertionError(_error_message()) from cause
AssertionError: Expected call: some_method3(a=1, b=5)
Actual call: some_method3(a=1, b=4)

代码实例

这里是不同环境下模拟代码的方法，它们都采用下面这些基础代码：

<span id="code_example_class_demo"></span>
类代码风格的基本代码

import os

# 面向对象开发中，往往需要载入其它模块，这个 ExternalClass 就用于模拟其它开发人员写的模块，我们既不知道它怎么做，也不知道做的对不对，只想模拟调用这个方法之后的结果
class ExternalClass(object):
    def __init__(self):
        self.external_attrib_a = None
        self.external_attrib_b = None

    def external_method_a(self):
        pass

class MyClass(object):
    def __init__(self):
        self.external = ExternalClass()

        self.attrib_a = None
        self.attrib_b = None

    def method_a(self):
        return self.external.external_method_a()

<span id="code_example_function_demo"></span>

过程风格的基本代码

Mock 类中的属性

这个例子中准备测试 MyClass 中的 method_a 方法， method_a 则实例化 ExternalClass 类，并调用它的 external_method_a 方法。

我们不在乎 external_method_a 怎么干的，就想模拟它的返回值。这就要用到 mock.PropertyMock 方法

class TestMyClass(object):
    # 因为要调用外部类，所以这里先把这个类实例化，在示例源码中也可以看到 MyClass 的 __init__ 方法中也是实例化了外部类的
    def setUp(self):
        self.ins = MyClass()

    # mock.PropertyMock 专用于模拟类中的属性 (不是方法，方法用 object)，关键就是不管实际代码中怎么实例化，或者实例化成什么名字，我们始终只模拟那个外部类
    def test_method_a(self):
        ExternalClass.external_attrib_a = mock.PropertyMock(return_value="hello")
        self.assertEqual(self.ins.method_a(), "hello")

Mock 文件的读写

代码中有时候要用 open 读写文件，下面的例子用于文件读写。关键就是 mock_open 操作。

下面的代码先用 open 打开文件，然后在里面用 read, write 操作文件，那么测试代码中就 mock "builtins.open" ，然后模拟 read 和 write 动作。

# 这是写文件的代码
def operate_file(file_name, content):
    with open(os.path.expanduser(filename), "wb") as fid:
        fid.write(content)

# 这里是测试代码
# mock.patch 用于模拟系统的 open 方法
from unittest import mock
@mock.patch("builtins.open", read_data="data")
def test_operate_file(mock_open):
    mock_open.read.return_value = True
    mock_open.write.return_value = True

# 上面是在 Python2 有效的代码，在 Python3 中 mock 建立了一个 mock_open 方法用来直接模拟，不需要用装饰器了，直接在函数内部这么写
def test_operate_file():
    with mock.patch("builtins.open", mock.mock_open(read_data=conf_lines)) as mock_open:
        ......

但是如果操作文件是在类里面，而且直接 open 文件以后用 for 循环文件句柄，没有 read, write 动作应该怎么做呢？看下面的例子：

# 这里是直接操作文件的代码
class MyClass(object):
    def handle_file(filename):
        with open(filename, "rt") as fid:
            for line in fid:
                ...


# 模拟上面的文件操作关键是要模拟 __iter__ 生成器
from unittest import TestCase, mock
class TestParser(TestCase):
    def setUp(self):
        self.ins = MyClass()

    def test_handle_file(self):
        with mock.patch("builtins.open") as mock_open:
            mock_open.return_value.__enter__ = mock_open
            mock_open.return_value.__iter__ = mock.Mock(return_value=iter(file_lines.splitlines()))

            # filename 的参数任意，反正没什么用，上面已经模拟的读文件的结果
            response = self.ins.handle_file(filename="fake_filename")

使用 with 语法

上面介绍过适用于函数内部直接用 mock_func = mock.patch() 模拟的方式，也介绍过在函数或方法上用装饰器 @mock.patch() 的方式模拟。除此之外还可以用 with 语句模拟，比如下面几段代码的功能是相同的：

# 函数内部直接模拟
import os
from unittest import mock

def function():
    mock_func = mock.patch("os.path.isfile", return_value=True)


# 使用装饰器
@mock.patch("os.path.isfile")
def function(mock_os_isfile):
    mock_os_isfile.return_value = True


# 使用 with 语句
def function():
    with mock.patch("os.path.isfile") as mock_os_isfile:
        mock_os_isfile.return_value = True

如果同时模拟多个模块或方法，那么多个 mock 之间用斜杠分隔，就像这样：

def test_run(self):
    """UT Case"""
    with mock.patch.object(PolicyLookup, "compare_zone", return_value=None) as mock_compare_zone, \
         mock.patch.object(PolicyLookup, "write_data_to_database", return_value=None) as mock_write_data_to_database:

Mock 实例

面向对象过程中可能有需要 mock 实例的情况，比如下面代码中有一个设备对象，设备有 login 方法，现在要测试的是类 OperateDevice 中的 login_device 方法，这时就涉及要模拟 Device 类中的 login 方法的问题了。可是在 login_device 中用的是 Device 类的一个实例啊，怎么把实例和类关联起来呢？

Class Device(object):
    def __init__(self):
        pass

    def login(self):
        pass


Class OperateDevice(object):
    def __init__(self):
        self.ins = Device()

    def login_device(self):
        self.ins.login()


Class TestOperateDevice(TestCase):
    def __init__(self):
        self.ins_operate_device = OperateDevice()

    def test_login_device(self):
        dev_obj = mock.Mock()
        dev_obj.login_device = mock.Mock()
        dev_obj.login_device.return_value = True

        self.ins_operate_device.login_device()

上面的例子测试的是 OperateDevice ，在里面的实例是类 Device，反正不管怎样，我们要模拟 login_device ，那么直接用 mock.Mock() 模拟一个类实例，然后再模拟一个方法并设置方法的值即可

原文链接：https://www.jianshu.com/p/55e5a6863c3f