让unittest框架在fail时自动截图的装饰器

---

在使用unittest测试框架的时候，我们在判断测试结果的时候尝尝会使用断言。在使用appium框架时，我们可以直接使用assert语句进行判断。而我们在测试时，有时候会想要让代码在用例fail的时候自动生成截图。但是使用assert方法进行判断时，fail的用例只直接不执行之后的代码的，所以我们只能通过使用try…exception语句来让程序进行截图操作。

            try:
                self.assertEqual(u'xxx','xxx','对比失败')
            except AssertionError:
            #调用ADB命令进行截图操作
               path = os.getcwd() + "/screenshot"
                 timestamp = time.strftime('%Y-%m-%d-%H-%M-%S', time.localtime(time.time()))
                 os.popen("adb wait-for-device")
                 os.popen("adb shell screencap -p /data/local/tmp/tmp.png")
                 if not os.path.isdir(os.getcwd() + "/screenshot"):
                     os.makedirs(path)
                 os.popen("adb pull /data/local/tmp/tmp.png " + path + "/" + timestamp + ".png")
                 os.popen("adb shell rm /data/local/tmp/tmp.png") 
            raise AssertionError,msg

但是这么实现之后，又觉得每一条assert语句都需要写一条try..exception.
是不是太麻烦了。于是就想到了之前看过装饰器。于是百度了解了一下原理。就开始尝试在unittest框架中给assert加上装饰器。来实现try..exception.

先找到assert语句块的位置。在C:\Python27\Lib\unittest目录里的case.py。打开文件。我进行了如下修改。

# -*- coding:utf-8 -*-
# """Test case implementation"""

import collections
import sys
import functools
import difflib
import pprint
import re
import types
import warnings
#这里引入了两个新模块
import time 
import os

然后在class TestCase(object):中

class TestCase(object):
    """A class whose instances are single test cases.

    By default, the test code itself should be placed in a method named
    'runTest'.

    If the fixture may be used for many test cases, create as
    many test methods as are needed. When instantiating such a TestCase
    subclass, specify in the constructor arguments the name of the test method
    that the instance is to execute.

    Test authors should subclass TestCase for their own tests. Construction
    and deconstruction of the test's environment ('fixture') can be
    implemented by overriding the 'setUp' and 'tearDown' methods respectively.

在类里添加一个ADD装饰器。
装饰器简单的理解就给函数一个包装。这个装饰器叫做ADD。之后的函数只要加上@add。就是加上了这个装饰器。那么在这个函数运行的时候，就会加上我在add里写的功能。实际上，就是给函数加上了一个外壳，增加了一些功能。
这个add函数，说了很简单，当接收到一个抛出来的AssertionError的时候，就是调用adb命令，把屏幕截取下来，放到一个文件夹下。

############这里定义了装饰器add##########
    def add(func):
        def wrapper(self, first, second, msg=None):
            try:
                func(self, first, second, msg=None)
            except AssertionError:#等待AssertionError
                 path = os.getcwd() + "/screenshot"
                 timestamp = time.strftime('%Y-%m-%d-%H-%M-%S', time.localtime(time.time()))
                 os.popen("adb wait-for-device")
                 os.popen("adb shell screencap -p /data/local/tmp/tmp.png")
                 if not os.path.isdir(os.getcwd() + "/screenshot"):
                     os.makedirs(path)
                 os.popen("adb pull /data/local/tmp/tmp.png " + path + "/" + timestamp + ".png")
                 os.popen("adb shell rm /data/local/tmp/tmp.png")
            raise AssertionError,msg#抛出AssertionError和信息
        return wrapper

    然后给以下方法加上装饰器

########################给方法加上装饰器
    @add
    def _baseAssertEqual(self, first, second, msg=None):
        """The default assertEqual implementation, not type specific."""
        if not first == second:
            standardMsg = '%s != %s' % (safe_repr(first), safe_repr(second))
            msg = self._formatMessage(msg, standardMsg)
            raise self.failureException(msg)
########################给方法加上装饰器
    @add
    def assertEqual(self, first, second, msg=None):
        """Fail if the two objects are unequal as determined by the '=='
           operator.
        """
        assertion_func = self._getAssertEqualityFunc(first, second)
        assertion_func(first, second, msg=msg)


########################给方法加上装饰器
    @add
    def assertNotEqual(self, first, second, msg=None):
        """Fail if the two objects are equal as determined by the '!='
           operator.
        """
        if not first != second:
            msg = self._formatMessage(msg, '%s == %s' % (safe_repr(first),
                                                          safe_repr(second)))
            raise self.failureException(msg)

只给assertEqual和assertNotEqual加了装饰器，所以，我们如果想要assert结束后能产生截图，那我们就只能用这两个方法了。
但是unittest框架中其实提供了很多方法。你如果相用的话，也可以把他们都加上这个装饰器。
但是我理解上，这些方法本质上都是能够用assertEqual和assertNotEqual代替的，所以就只改了这两个，毕竟这样改动比较简单嘛。。。

然后就可以实现了。

case.py全文：

# -*- coding:utf-8 -*-
# """Test case implementation"""

import collections
import sys
import functools
import difflib
import pprint
import re
import types
import warnings
import time

import os

from . import result
from .util import (
    strclass, safe_repr, unorderable_list_difference,
    _count_diff_all_purpose, _count_diff_hashable
)


__unittest = True


DIFF_OMITTED = ('\nDiff is %s characters long. '
                 'Set self.maxDiff to None to see it.')

class SkipTest(Exception):
    """
    Raise this exception in a test to skip it.

    Usually you can use TestCase.skipTest() or one of the skipping decorators
    instead of raising this directly.
    """
    pass

class _ExpectedFailure(Exception):
    """
    Raise this when a test is expected to fail.

    This is an implementation detail.
    """

    def __init__(self, exc_info):
        super(_ExpectedFailure, self).__init__()
        self.exc_info = exc_info

class _UnexpectedSuccess(Exception):
    """
    The test was supposed to fail, but it didn't!
    """
    pass

def _id(obj):
    return obj

def Myskip(func):
    def RebackTest(self):
        if self._resultForDoCleanups.failures or self._resultForDoCleanups.errors:
            raise unittest.SkipTest("{} do not excute because {} is failed".format(func.__name__,self._resultForDoCleanups.failures[0][0]._testMethodName))
        func(self)
    return  RebackTest

def skip(reason):
    """
    Unconditionally skip a test.
    """
    def decorator(test_item):
        if not isinstance(test_item, (type, types.ClassType)):
            @functools.wraps(test_item)
            def skip_wrapper(*args, **kwargs):
                raise SkipTest(reason)
            test_item = skip_wrapper

        test_item.__unittest_skip__ = True
        test_item.__unittest_skip_why__ = reason
        return test_item
    return decorator

def skipIf(condition, reason):
    """
    Skip a test if the condition is true.
    """
    if condition:
        return skip(reason)
    return _id

def skipUnless(condition, reason):
    """
    Skip a test unless the condition is true.
    """
    if not condition:
        return skip(reason)
    return _id


def expectedFailure(func):
    @functools.wraps(func)
    def wrapper(*args, **kwargs):
        try:
            func(*args, **kwargs)
        except Exception:
            raise _ExpectedFailure(sys.exc_info())
        raise _UnexpectedSuccess
    return wrapper


class _AssertRaisesContext(object):
    """A context manager used to implement TestCase.assertRaises* methods."""

    def __init__(self, expected, test_case, expected_regexp=None):
        self.expected = expected
        self.failureException = test_case.failureException
        self.expected_regexp = expected_regexp

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_value, tb):
        if exc_type is None:
            try:
                exc_name = self.expected.__name__
            except AttributeError:
                exc_name = str(self.expected)
            raise self.failureException(
                "{0} not raised".format(exc_name))
        if not issubclass(exc_type, self.expected):
            # let unexpected exceptions pass through
            return False
        self.exception = exc_value # store for later retrieval
        if self.expected_regexp is None:
            return True

        expected_regexp = self.expected_regexp
        if not expected_regexp.search(str(exc_value)):
            raise self.failureException('"%s" does not match "%s"' %
                     (expe