功能介紹

好的編碼習(xí)慣都應(yīng)該為每一行代碼做覆蓋測試，但有些時候代碼處理的是從網(wǎng)絡(luò)上獲取的內(nèi)容，或者設(shè)備的返回，比如獲取交換機路由器的運行結(jié)果，或者從網(wǎng)絡(luò)上獲取頁面等等。這些動作要么需要聯(lián)網(wǎng)，要么需要設(shè)備，但實際上我們只是想測試代碼正確性而已，注重的是對返回的內(nèi)容的處理而不必非要有實際設(shè)備。

mock 模塊用于在單元測試中模擬其它代碼的結(jié)果，比如某個函數(shù)需要調(diào)用其他函數(shù)，這個時候我們可以模擬這個第三方函數(shù)的結(jié)果來略過實際調(diào)用它，不光可以節(jié)省時間，也可以避免因為第三方函數(shù)出錯而影響自己的代碼，甚至可以很輕松的模擬難以出現(xiàn)的各種情況。

也正是因為這個模塊是如此好用，在 Python2 中還需要單獨安裝 mock 模塊，而 Python3.3 開始這個模塊就被放入標(biāo)準(zhǔn)模塊了，名叫 unittest.mock

使用思路和實例

在概念上， mock 用于模擬函數(shù)的返回，比如你有一個函數(shù)調(diào)用了另一個函數(shù)，而另一個函數(shù)的代碼本身不是你寫的，或者不需要在當(dāng)前單元測試中測試，你只是希望拿到另一個函數(shù)返回的結(jié)果，這個時候就可以用 mock 來模擬那個函數(shù)來略過各種中間過程而直接得到結(jié)果。比如下面這樣的代碼結(jié)構(gòu)：

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

上面的 test_search_flow_session 是寫在單元測試腳本中的測試案例，用來測試在另一個源代碼文件中的 search_flow_session 函數(shù)。而 search_flow_session 要調(diào)用另 2 個其它文件中的函數(shù) send_shell_cmd 和 get_all_flow_session 來完成功能。恰恰麻煩的是這 2 個函數(shù)其中一個需要一臺 PC 機來執(zhí)行 linux 命令，另一個需要一臺昂貴的設(shè)備來獲取設(shè)備上的狀態(tài)和返回，更別說創(chuàng)建拓?fù)浜突謴?fù)測試環(huán)境的工作，僅僅為了檢查 search_flow_session 中的某些代碼而付出這樣的代價完全不值。

但是應(yīng)該怎么用 mock 模擬，或者怎么把 mock 注入到你自己的函數(shù)中卻是一個很傷腦筋的問題，不同的代碼風(fēng)格很容易把你帶進坑里，比如要調(diào)用的其他函數(shù)使用 OOP 方式寫的，你會想難道我還得先實例化？或者我的函數(shù)是面向?qū)ο蟮?，調(diào)用的卻是面向過程的，怎么辦？在我剛剛開始接觸 mock 的時候，這些概念和行為真是把我折磨的夠嗆。寫多了之后才慢慢感覺到了下面幾個規(guī)則：

1. 不用管自己的函數(shù)怎么寫， mock 只用來模擬別人的模塊，不管是面向過程還是面向?qū)ο蠖疾挥眠^多考慮，只考慮你的代碼中調(diào)用了哪些外部函數(shù)或者方法，這意味著你要 mock 多少東西
2. 如果調(diào)用的外部代碼是面向過程的風(fēng)格，也就是一個一個函數(shù)，那么用 mock.patch ；面向?qū)ο箫L(fēng)格，比如你調(diào)用的只是一個類中的某個方法則用 mock.patch.object ?，F(xiàn)在看到什么 mock.patch , mock.patch.object 可能你不理解，沒事，先放下，到后面會專門說

mock 概念很繞，但是真正用到的接口并不多。也是，模擬函數(shù)或者方法行為而已，又能有幾種接口呢……大致說來我們能接觸到的也就是這么幾個：

Mock

mock 是最初，也是最基本的一個函數(shù)，它的任務(wù)就是模擬某個模塊的函數(shù)。

patch - 補丁方式模擬

有些函數(shù)可能不屬于你，你也不在意它的內(nèi)部實現(xiàn)而只是想調(diào)用這個函數(shù)然后得到結(jié)果而已，這種時候就可以用 patch 方式來模擬。

比如一個模塊 linux_tool.py 里面有多個函數(shù)，其中 send_shell_cmd 是其他人寫的。它具體怎么做我不在乎，只知道它向 Linux PC 發(fā)命令然后將命令的結(jié)果返回給我?，F(xiàn)在我寫了一個函數(shù) check_cmd_response 檢查返回結(jié)果，然后對 check_cmd_response 做單元測試。因為 send_shell_cmd 函數(shù)需要一個真實的 PC ，這需要設(shè)備且每次返回還可能與預(yù)期不符，比如設(shè)備無法連接，想檢查的東西忘記配置所以取不回來等等，這些都會干擾我自己函數(shù)的行為，而且問題和自己函數(shù)無關(guān)，這種時候就可以用 mock 模擬 send_shell_cmd 函數(shù)而且把預(yù)期返回寫到這個模擬過程中，保證每次都會正確處理。當(dāng)然有人說可能的確有錯誤情況啊，這也是你應(yīng)該要處理的，或者有多種返回啊……沒錯，所以可以多寫幾個測試案例把這些情況都模擬一遍嘛。

面向過程代碼風(fēng)格

下面是完整的模擬代碼，首先是 linux_tool.py 文件，里面 2 個函數(shù)， send_shell_cmd 直接返回一個字符串，注意在現(xiàn)實中這是一個完整函數(shù)會連接設(shè)備并獲取返回的。另一個就是自己寫的函數(shù)了，中間的代碼都去掉，但是整體來說我希望獲取未來使用 mock 模擬的函數(shù)所返回的內(nèi)容

#!/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 的用法，它是一個裝飾器，需要把你想模擬的函數(shù)寫在里面，然后在后面的單元測試案例中為它賦一個具體實例，再用 return_value 來指定模擬的這個函數(shù)希望返回的結(jié)果就可以了，后面就是正常單元測試代碼。

#!/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 個文件中的，第一個關(guān)卡往往就是怎么把這 2 個文件有機結(jié)合起來。這里的關(guān)鍵就是：源代碼該怎么寫就怎么寫，不需要考慮為 mock 留下什么接口之類的東西。

2. 單元測試文件中，首先寫單元測試代碼，就和正常的一樣，最開始的時候只需要 import mock 模塊即可。

3. 判斷要測試的函數(shù)中是否用了其他函數(shù)，有可能使用了多個外部函數(shù)，那么就判斷哪個函數(shù)適合 mock ，哪些不需要，一般像浪費時間的，結(jié)果不定的，需要其他設(shè)備的函數(shù)最好都 mock ，其它一些功能函數(shù)可用可不用。

4. 確定了哪些外部函數(shù)要 mock 就用 patch 語句將它們列出來，每個 patch 是一個函數(shù)，而且要確定這些外部函數(shù)都在文件頭部用 import 語句載入到內(nèi)存了，因為 mock 模塊是通過替換內(nèi)存中的函數(shù)微代碼來實現(xiàn)功能的。

5. 如果 patch 多個外部函數(shù)，那么調(diào)用遵循自下而上的規(guī)則，比如：

   @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()))
   

如果函數(shù)是在其它文件中實現(xiàn)的，那么 mock 的方式又有不同：

# run_multiple 是在另一個文件 utils.py 中實現(xiàn)的
def run_multiple():
    pass

# 但是在 tool.py 文件中調(diào)用了這個模塊
from utils import run_multiple

def tool():
    run_multiple()

# test_tool.py 測試的時候就不能 mock 原始實現(xiàn)的路徑，而是使用的路徑
import unittest2 as unittest
import mock

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

上面的關(guān)鍵就是 mock.patch 的路徑必須是 "tool.run_multiple" ，這是使用 run_multiple 函數(shù)的路徑，而不是實現(xiàn)這個函數(shù)的路徑 "utils.run_multiple"

面向?qū)ο蟠a風(fēng)格

如果你的代碼風(fēng)格是面向?qū)ο蟮哪?？也可以，?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)

面向?qū)ο蟮?mock 和面向過程的很相似，唯一就是把 mock.patch 替換成 mock.patch.object ，并且在里面列出類實例和方法名。仔細觀察，是類的實例 (不是字符串) 和方法名 (是字符串的方法名而不是方法對象)

side_effect

side_effect 是 mock 中角色比較復(fù)雜的方法，它有好幾種用法

模擬同一個函數(shù)被多次調(diào)用

如果要多次調(diào)用相同函數(shù)并獲取返回，比如有一個外部方法叫 linux_tool.send_shell_cmd 用來執(zhí)行命令并返回命令中間的輸出，利用這個函數(shù)我又寫了一個自己的方法用來建立 vsftpd 服務(wù)器，其中多次調(diào)用外部方法來創(chuàng)建備份文件，建立配置文件，重啟服務(wù)，檢查服務(wù)狀態(tài)等等。或者某個命令在一個循環(huán)中被調(diào)用，循環(huán)次數(shù)也可能是不定的。上面的例子都只是模擬了一次，那么模擬多次怎么辦？

答案就是使用 side_effect ，比如下面的例子中在方法 start_ftp_service 中調(diào)用了 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()

如果某個命令在循環(huán)中被調(diào)用，滿足判斷結(jié)果才會跳出循環(huán)，那么也要用 side_effect 來模擬循環(huán)中的每次結(jié)果，一定數(shù)清楚具體的循環(huán)次數(shù)或者精心設(shè)計返回，否則執(zhí)行會出錯。

模擬異常

用上面模擬同一個函數(shù)多次被調(diào)用的實例為例，如果希望主動引發(fā)異常，比如 Exception 那么可以這樣：

mock_send_shell_cmd.side_effect = Exception("Raise Exception")

所有 raise 語句可以引發(fā)的異常都可以用 side_effect 引發(fā)

模擬對象中的屬性

有些時候要模擬的不是其它類中的方法，而是屬性，比如下面這個類里面有一個屬性 before ，一個方法 spawnu ，方法的模擬很簡單在上面已經(jīng)有說明，但 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 的子類，區(qū)別就是 MagicMock 預(yù)置了其它 MagicMethod ，所謂 MagicMethod 在 Python 中表現(xiàn)為雙下劃線包圍的方法，比如最熟悉的 init 或者 str 之類的。 mock.Mock 默認(rèn)沒有實現(xiàn)這些方法，如果想測試這些方法的行為就得自己寫，而 MagicMock 默認(rèn)預(yù)置了這些行為，這樣像自增自減，列表的循環(huán)，計算符號的重載等 MagicMethod 就在 MagicMock 中內(nèi)置了，如果不考慮這些那么 MagicMock 和 Mock 行為是一樣的

一般情況下模擬都用 MagicMock ，因為這個模擬出來的行為更類似于我們預(yù)期

精準(zhǔn)模擬第三方函數(shù)

自己寫的模塊大多數(shù)時候都需要調(diào)用其它函數(shù) (比如大多數(shù)模塊都會用的 os 或者 sys 模塊) ，如何模擬這些第三方函數(shù)呢？可以看一個例子：

# 功能模塊， 模塊名 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 ，這個模塊很可能在很多模塊中都被載入和調(diào)用過，如果源碼文件特別多的話可能 os 這個模塊會到處都是，而測試代碼中如果直接模擬 os 模塊的話很可能多個 test_ 源文件會互相影響。最好的辦法就是對每個源文件的第三方模塊精準(zhǔn)模擬

在 demo.py 文件中調(diào)用了 os.path.isfile 和 os.remove 方法，如何精準(zhǔn)模擬呢？上面的例子中用 mock.patch("demo.os...") 的方式就可以做到

內(nèi)建的其他方法

called

一旦 mock 被創(chuàng)建，比如上面用 patch 模擬的 mock_send_shell_cmd ，或者用 MagicMock 模擬的 mock_func ，都可以用 called() 方法來檢查自己究竟有沒有被調(diào)用，比如：

mock_send_shell_cmd.called
>> True

call_count

返回模擬的函數(shù)或方法被調(diào)用了幾次：

mock_send_shell_cmd.call_count
>> 2

call_args

返回 mock 的東西在調(diào)用時傳入的具體參數(shù)

>>> 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 的方法被多次調(diào)用的情況，會返回一個列表，列表中是每次被調(diào)用時的參數(shù)

assert_called_with

有時候我們不光想確認(rèn)自己 mock 的東西有沒有被調(diào)用，還想確認(rèn)調(diào)用時傳入的參數(shù)是不是正確的，就可以用 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)

代碼實例

這里是不同環(huán)境下模擬代碼的方法，它們都采用下面這些基礎(chǔ)代碼：

<span id="code_example_class_demo"></span>
類代碼風(fēng)格的基本代碼

import os

# 面向?qū)ο箝_發(fā)中，往往需要載入其它模塊，這個 ExternalClass 就用于模擬其它開發(fā)人員寫的模塊，我們既不知道它怎么做，也不知道做的對不對，只想模擬調(diào)用這個方法之后的結(jié)果
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>

過程風(fēng)格的基本代碼

Mock 類中的屬性

這個例子中準(zhǔn)備測試 MyClass 中的 method_a 方法， method_a 則實例化 ExternalClass 類，并調(diào)用它的 external_method_a 方法。

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

class TestMyClass(object):
    # 因為要調(diào)用外部類，所以這里先把這個類實例化，在示例源碼中也可以看到 MyClass 的 __init__ 方法中也是實例化了外部類的
    def setUp(self):
        self.ins = MyClass()

    # mock.PropertyMock 專用于模擬類中的屬性 (不是方法，方法用 object)，關(guān)鍵就是不管實際代碼中怎么實例化，或者實例化成什么名字，我們始終只模擬那個外部類
    def test_method_a(self):
        ExternalClass.external_attrib_a = mock.PropertyMock(return_value="hello")
        self.assertEqual(self.ins.method_a(), "hello")

Mock 文件的讀寫

代碼中有時候要用 open 讀寫文件，下面的例子用于文件讀寫。關(guān)鍵就是 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 用于模擬系統(tǒng)的 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 方法用來直接模擬，不需要用裝飾器了，直接在函數(shù)內(nèi)部這么寫
def test_operate_file():
    with mock.patch("builtins.open", mock.mock_open(read_data=conf_lines)) as mock_open:
        ......

但是如果操作文件是在類里面，而且直接 open 文件以后用 for 循環(huán)文件句柄，沒有 read, write 動作應(yīng)該怎么做呢？看下面的例子：

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


# 模擬上面的文件操作關(guān)鍵是要模擬 __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 的參數(shù)任意，反正沒什么用，上面已經(jīng)模擬的讀文件的結(jié)果
            response = self.ins.handle_file(filename="fake_filename")

使用 with 語法

上面介紹過適用于函數(shù)內(nèi)部直接用 mock_func = mock.patch() 模擬的方式，也介紹過在函數(shù)或方法上用裝飾器 @mock.patch() 的方式模擬。除此之外還可以用 with 語句模擬，比如下面幾段代碼的功能是相同的：

# 函數(shù)內(nèi)部直接模擬
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 實例

面向?qū)ο筮^程中可能有需要 mock 實例的情況，比如下面代碼中有一個設(shè)備對象，設(shè)備有 login 方法，現(xiàn)在要測試的是類 OperateDevice 中的 login_device 方法，這時就涉及要模擬 Device 類中的 login 方法的問題了?？墒窃?login_device 中用的是 Device 類的一個實例啊，怎么把實例和類關(guān)聯(lián)起來呢？

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() 模擬一個類實例，然后再模擬一個方法并設(shè)置方法的值即可