unittest+selenium测试失败截图 最佳实现

网上许多类似的装饰器是不好用的，本人设计实现的装饰器可以不带参数，也可以带参数，可同时用于func与class。
只需要满足一个基本条件，那就是测试类需要包含一个名为driver的WebDriver的类属性或实例属性。

class ClassA(unittest.TestCase):
    driver = None  # driver是一个WebDriver实例

    @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是函数式装饰器，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"):
    """
    一个函数装饰器，用于unittest执行测试selenium的用例出现失败后，自动截图

import unittest
import selenium.webdriver as webdriver


# example_1:
@screen
class ClassA(unittest.TestCase):
    driver = None  # driver是一个WebDriver实例

    @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实例

    @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: 当装饰class时，可以用于标记哪些测试方法会被自动装饰
    :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对应测试类中的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实例

    @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实例

    @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对应测试类中的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