源碼教學：教你 30 行代碼實作 ddt 子產品

前言

用 python 做過自動化的小夥伴，大多數都應該使用過 ddt 這個子產品，不可否認 ddt 這個子產品确實挺好用，可以自動根據用例資料，來生成測試用例，能夠很友善的将測試資料和測試用例執行的邏輯進行分離。接下來就帶大家一起自己，手把手撸出一個 ddt。

1、DDT 的實作原理

首先我們來看一下 ddt 的基本使用：

ddt 在使用時非常簡潔，也就是兩個裝飾器，@ddt 這個裝飾器裝飾測試類，@data 這個裝飾器裝飾器用例方法并傳入測試資料。這兩個裝飾器實作的效果就是根據傳入的用例資料自動生成用例。具體是怎麼實作的呢？其實實作的思路也特别的簡單，也就兩個步驟：

整理了一份大廠軟體測試面試寶典pdf

第一步：把傳進來的用例資料儲存起來

第二步：周遊用例資料，每周遊一條資料 就動态的給測試類添加一個用例方法。

ddt 中的兩個裝飾器其實實作的就是這麼兩個步驟：

@data:做的是第一步将傳入測試資料儲存起來;

@ddt 做的是第二步，周遊用例資料，給測試類動态添加用例方法。

2、data 裝飾器的實作

前面我們說到 data 這個裝飾器，做的事情是将用例資料儲存起來。那麼如何儲存呢？其實最簡單的方式就是 儲存被裝飾的這個用例方法的屬性。接下來我們來具體實作：

先看一個 ddt 使用的案例

@ddt
class TestLogin(unittest.TestCase):

    @data(11,22)
    def test_login(self, item):
        pass
           

了解過裝飾器裝飾器原理的小夥伴，應該都知道上面@data(11,22) 這行代碼執行的效果等同于

test_login = data(11,22)(test_login)
           

接下來我們來分析一下上面這行代碼，首先是調用 data 這個裝飾器函數，把用例資料 11,22 當成參數傳入進去，然後傳回一個可調用對象(函數)，再次調用傳回的函數并把用例方法傳入進去。明确了調用的流程那麼我們就可以結合之前的需求去定義 data 這個裝飾器函數了。具體實作如下：

def data(*args):
    def wrapper(func):
        setattr(func, "PARAMS", args)
        return func
    return wrapper
           

代碼解讀：

前面的案例在使用 data 時，執行的 test_login = data(11,22)(test_login) 先調用 data 傳入的 11,22 通過不定長參數 args 接收，然後傳回嵌套的函數 wrapper 然後調用傳回的 wrapper 函數，傳入被裝飾的 test_login 方法 在 wrapper 函數中我們把用例資料儲存為 test_login 這個方法的 PARAMS 屬性，再把 test_login 傳回 到此為止，data 這個裝飾器我們就實作用例資料的儲存

3、ddt 裝飾器的實作

通過 data 這個裝飾器我們實作了用例資料儲存之後，我們接下來實作 ddt 這個裝飾器，根據用例資料生成測試用例。前面的案例 @ddt 裝飾測試類的時候，實際上執行的效果等同于下面的代碼

TestLogin = ddt(TestLogin)
           

這行代碼就是把被裝飾器的類傳入到 ddt 這個裝飾器函數中，再把傳回值指派給 TestLogin。之前我們分析的時候說了 ddt 這個裝飾器做的事情是周遊用例資料，動态的給測試類添加用例方法，接下來我們就來實作 ddt 這個裝飾器内部的邏輯。

def ddt(cls):
    for name, func in list(cls.__dict__.items()):
        if hasattr(func, "PARAMS"):
            for index, case_data in enumerate(getattr(func, "PARAMS")):
                new_test_name ="{}_{}".format(name,index)
                setattr(cls, new_test_name, func)
            else:
                delattr(cls, name)
    return cls
           

ddt 函數内部邏輯說明： 1、調用 ddt 這個函數時會把測試類當成參數傳入進來， 2、然後通過 cls.__dict__ 擷取測試的所有屬性和方法，進行周遊 3、判斷變量出來的屬性或方法 有沒有 PARAMS 這個屬性， 4、如果有，則說明這個方法用 data 裝飾器裝飾過并傳入了用例資料。 5、通過 getattr(func, "PARAMS")擷取所有的用例資料，進行周遊。 6、每周遊出來一組用例資料，生産一個用例方法名， 再動态的給測試類添加一個用例方法。 7、周遊完所有用例資料之後，删除測試類原來定義的測試方法 8、最後傳回測試類

當目前為止 ddt 和 data 這兩個裝飾器函數的基本功能實作了，可以自動根據用例資料生成測試用例了，接下來我們寫個測試類來檢查一下

# 定義裝飾器函數data
def data(*args):
    def wrapper(func):
        setattr(func, "PARAMS", args)
        return func

    return wrapper

# 定義裝飾器函數ddt
def ddt(cls):
    for name, func in list(cls.__dict__.items()):
        if hasattr(func, "PARAMS"):
            for index, case_data in enumerate(getattr(func, "PARAMS")):
                new_test_name = "{}_{}".format(name, index)
                setattr(cls, new_test_name, func)
            else:
                delattr(cls, name)
    return cls

import unittest

# 編寫測試類
@ddt
class TestDome(unittest.TestCase):
    @data(11, 22, 33, 44)
    def test_demo(self):
        pass
           

運作上述用例，我們就會發現執行了四條用例，根據用例資料生成用例的功能就已經實作了。

4、解決用例參數傳遞的問題

雖然上面基本的功能已經實作了，但是還存在一個問題。用例的資料沒有傳遞到用例方法中。那麼用例資料傳遞怎麼實作了，我們可以通過一個閉包函數對用例方法進行修，進而實作在調用用例方法的時候，把用例測試當成參數傳遞進去。修改原有用例方法的函數代碼如下

from functools import wraps

def update_test_func(test_func,case_data):
    @wraps(test_func)
    def wrapper(self):
        return test_func(self, case_data)
    return wrapper
           

上面我們定義了一個叫做 update_test_func 的閉包函數 閉包函數接收兩個參數：test_func(接收用例方法),case_data(接收用例資料) 閉包函數傳回一個嵌套函數，嵌套函數内部調用原來的用例方法，并傳入測試資料 嵌套函數在定義時，使用了 functools 子產品中的裝飾器 wraps 來裝飾，它可以讓 wrapper 這個嵌套函數具有 test_func 這個用例函數的相關屬性。

下面我們回到前面寫的 ddt 這個函數中，在給測試類添加用例之前，調用 update_test_func 方法對用例方法進行修改。

def ddt(cls):
    for name, func in list(cls.__dict__.items()):
        if hasattr(func, "PARAMS"):
            for index, case_data in enumerate(getattr(func, "PARAMS")):
                # 生成一個用例方法名
                new_test_name = "{}_{}".format(name, index)
                # 修改原有的測試方法，設定用例資料為測試方法的參數
                test_func = update_test_func(func,case_data)
                setattr(cls, new_test_name, test_func)
            else:
                delattr(cls, name)
    return cls
           

通過加上這一步之後，我們在測試類中 動态給測試類添加的測試方法，其實指向的全部是 update_test_func 裡面定義的 wrapper 函數，在執行測試用的時候實際上也是執行的 wrapper 函數，而在 wrapper 函數内部，我們調用了原來定義的測試方法，并将用例資料傳入了進去，到此為止 ddt 的功能我們就完全實作了。

下面是一個完整的案例，大家可以複制過去運作，也可以自己去寫一遍，還可以根據自己的一些需求進行自定義的擴充。

完整案例

from functools import wraps
import unittest
# --------ddt的實作--------
def data(*args):
    def wrapper(func):
        setattr(func, "PARAMS", args)
        return func
    return wrapper
def update_test_func(test_func, case_data):
    @wraps(test_func)
    def wrapper(self):
        return test_func(self, case_data)
    return wrapper
def ddt(cls):
    for name, func in list(cls.__dict__.items()):
        if hasattr(func, "PARAMS"):
            for index, case_data in enumerate(getattr(func, "PARAMS")):
                # 生成一個用例方法名
                new_test_name = "{}_{}".format(name, index)
                # 修改原有的測試方法，設定用例資料為測試方法的參數
                test_func = update_test_func(func, case_data)
                setattr(cls, new_test_name, test_func)
            else:
                delattr(cls, name)
    return cls
# --------測試用例編寫--------
@ddt
class TestDome(unittest.TestCase):
    @data(11, 22, 33, 44)
    def test_demo(self, data):
        assert data < 40
#---------用例執行-----------
unittest.main()