網(wǎng)上許多類似的裝飾器是不好用的，本人設(shè)計(jì)實(shí)現(xiàn)的裝飾器可以不帶參數(shù)，也可以帶參數(shù)，可同時用于func與class。
只需要滿足一個基本條件，那就是測試類需要包含一個名為driver的WebDriver的類屬性或?qū)嵗龑傩浴?/p>

class ClassA(unittest.TestCase):
    driver = None  # driver是一個WebDriver實(shí)例

    @classmethod
    def setUpClass(cls):
        cls.driver = webdriver.Chrome()

    def tearDown(self):
        self.driver.close()
        self.driver.quit()

    def test_001(self):
        self.driver.get("http://www.baidu.com")
        raise AttributeError

以下為裝飾器源碼，有兩個裝飾器源碼，screen是函數(shù)式裝飾器，Screen是類式裝飾器，他們的功能是完全一樣的，使用方法也相同。裝飾測試類時，可@screen(func_prefix="test")或@Screen(func_prefix="test")，這樣只針對前綴為"test"的function失敗時截圖。

# coding=utf-8
import functools
import inspect


def screen(target=None, func_prefix="test"):
    """
    一個函數(shù)裝飾器，用于unittest執(zhí)行測試selenium的用例出現(xiàn)失敗后，自動截圖

import unittest
import selenium.webdriver as webdriver


# example_1:
@screen
class ClassA(unittest.TestCase):
    driver = None  # driver是一個WebDriver實(shí)例

    @classmethod
    def setUpClass(cls):
        cls.driver = webdriver.Chrome()

    def tearDown(self):
        self.driver.close()
        self.driver.quit()

    def test_001(self):
        self.driver.get("http://www.baidu.com")
        raise AttributeError


# example_2:
class ClassB(unittest.TestCase):
    driver = None  # driver是一個WebDriver實(shí)例

    @classmethod
    def setUpClass(cls):
        cls.driver = webdriver.Chrome()

    def tearDown(self):
        self.driver.close()
        self.driver.quit()

    @screen
    def test_001(self):
        self.driver.get("http://www.baidu.com")
        raise AttributeError


    :param target: 被裝飾的對象，可以是class, function
    :param func_prefix: 當(dāng)裝飾class時，可以用于標(biāo)記哪些測試方法會被自動裝飾
    :return: wrapped class 或 wrapped function
    """

    def decorator(func_or_cls):
        if inspect.isfunction(func_or_cls):
            @functools.wraps(func_or_cls)
            def wrapper(*args, **kwargs):
                try:
                    return func_or_cls(*args, **kwargs)
                except Exception:  # 可以修改要捕獲的異常類型
                    args[0].driver.save_screenshot(str(args[0]) + ".png")  # args[0].driver對應(yīng)測試類中的driver
                    raise

            return wrapper
        elif inspect.isclass(func_or_cls):
            for name, func in list(func_or_cls.__dict__.items()):
                if inspect.isfunction(func) and name.startswith(func_prefix):
                    setattr(func_or_cls, name, decorator(func))
            return func_or_cls
        else:
            raise AttributeError

    if target:
        return decorator(target)
    else:
        return decorator


class Screen(object):
    """
    類裝飾器, 功能與screen一樣

import unittest
import selenium.webdriver as webdriver


# example_1:
@Screen
class ClassA(unittest.TestCase):
    driver = None  # driver是一個WebDriver實(shí)例

    @classmethod
    def setUpClass(cls):
        cls.driver = webdriver.Chrome()

    def tearDown(self):
        self.driver.close()
        self.driver.quit()

    def test_001(self):
        self.driver.get("http://www.baidu.com")
        raise AttributeError


# example_2:
class ClassB(unittest.TestCase):
    driver = None  # driver是一個WebDriver實(shí)例

    @classmethod
    def setUpClass(cls):
        cls.driver = webdriver.Chrome()

    def tearDown(self):
        self.driver.close()
        self.driver.quit()

    @Screen
    def test_001(self):
        self.driver.get("http://www.baidu.com")
        raise AttributeError
    """

    def __new__(cls, func_or_cls=None, func_prefix="test"):
        cls._prefix = func_prefix
        self = object.__new__(cls)
        if func_or_cls:
            return self(func_or_cls)
        else:
            return self

    def __init__(self, func_or_class=None, func_prefix="test"):
        pass

    def __call__(self, func_or_cls=None):
        if inspect.isfunction(func_or_cls):
            @functools.wraps(func_or_cls)
            def wrapper(*args, **kwargs):
                try:
                    return func_or_cls(*args, **kwargs)
                except Exception:  # 可以修改要捕獲的異常類型
                    args[0].driver.save_screenshot(str(args[0]) + ".png")  # args[0].driver對應(yīng)測試類中的driver
                    raise

            return wrapper
        elif inspect.isclass(func_or_cls):
            for name, func in list(func_or_cls.__dict__.items()):
                if inspect.isfunction(func) and name.startswith(self._prefix):
                    setattr(func_or_cls, name, self(func))
            return func_or_cls
        else:
            raise AttributeError