Python 天天美味(37) - 让python的unittest像gtest一样输出

Python自带的unittest已经很简单易用了，不过我一直不喜欢的是它的命令行输出，格式显得有点乱。而我比较喜欢的是gtest的命令行输出格式，用不同的颜色进行标识，整齐划一，非常明了。于是，我扩展一下Python的unittest模块，让它也能输出和gtest一样好看的命令行结果。

首先，我们先来看看unittest默认的命令行输出结果。这里，我先随便写两个测试案例，让其中一个通过，另外一个不通过，然后查看一下测试结果。

import unitest

class FooTest(unittest.TestCase):
    def setUp(self):
        self.a = 1

    def testPass(self):
        self.a = self.a + 1
        self.assertEqual(2, self.a)

    def testFail(self):
        self.a = self.a + 1
        self.assertEqual(3, self.a)

main函数，调用unittest自己的TextTestRunner：

if __name__=='__main__':
unittest.main()

输出的结果：

有点凌乱（当然，可能你并不觉得），好的，接下来开始实现一个自定义的TestRunner，让unittest输出和gtest一样，使用不同的颜色。

这个过程其实很简单，主要分为两个步骤：

1. 编写自定义的TestRunner类，执行其中的run方法，控制整个测试的过程和输出。参照unitest自己的TextTestRunner方法就好了。只是把输出部分做一些修改。

2. 编写自定义的TestResult类，继承自unittest中的TestResult类。重写其中几个方法，也都是修改输出内容的部分。

当然，我们还需要在命令行中输出不同的颜色。

如果你不是使用的Windows，可以参照：http://code.activestate.com/recipes/475116-using-terminfo-for-portable-color-output-cursor-co/

如果你使用的是Windows，其实只要调用一个Windows API就好了。这个API就是SetConsoleTextAttribute。见下面的代码：

## {{{ http://code.activestate.com/recipes/496901/ (r3)
# See http://msdn.microsoft.com/library/default.asp?url=/library/en-us/winprog/winprog/windows_api_reference.asp
# for information on Windows APIs.
STD_INPUT_HANDLE = -10
STD_OUTPUT_HANDLE= -11
STD_ERROR_HANDLE = -12

FOREGROUND_WHITE = 0x0007
FOREGROUND_BLUE = 0x01 # text color contains blue.
FOREGROUND_GREEN= 0x02 # text color contains green.
FOREGROUND_RED  = 0x04 # text color contains red.
FOREGROUND_INTENSITY = 0x08 # text color is intensified.
FOREGROUND_YELLOW = FOREGROUND_RED | FOREGROUND_GREEN

BACKGROUND_BLUE = 0x10 # background color contains blue.
BACKGROUND_GREEN= 0x20 # background color contains green.
BACKGROUND_RED  = 0x40 # background color contains red.
BACKGROUND_INTENSITY = 0x80 # background color is intensified.

import ctypes

std_out_handle = ctypes.windll.kernel32.GetStdHandle(STD_OUTPUT_HANDLE)

def set_color(color, handle=std_out_handle):
    """(color) -> BOOL
    Example: set_color(FOREGROUND_GREEN | FOREGROUND_INTENSITY)
    """
    bool = ctypes.windll.kernel32.SetConsoleTextAttribute(handle, color)
    return bool

TestRunner类

class MyTestRunner:
    def __init__(self, stream=sys.stderr, descriptions=1, verbosity=1):
        self.stream = _ColorWritelnDecorator(stream)
        self.descriptions = descriptions
        self.verbosity = verbosity

    def run(self, test):
        result = MyTestResult(self.stream, self.descriptions, self.verbosity)
        self.stream.yellow('Note: Your Unit Tests Starts')
        self.stream.writeln()
        startTime = time.time()
        test(result)
        stopTime = time.time()
        timeTaken = stopTime - startTime
        self.stream.green(result.separator2)
        run = result.testsRun
        self.stream.writeln("Ran %d test%s in %.3fs" %
                            (run, run != 1 and "s" or "", timeTaken))

        failed, errored = map(len, (result.failures, result.errors))

        self.stream.green("[  PASSED  ] %d tests" % (run - failed - errored))
        self.stream.writeln()

        if not result.wasSuccessful():
            errorsummary = ""
            if failed:
                self.stream.red("[  FAILED  ] %d tests, listed below:" % failed)
                self.stream.writeln()
                for failedtest, failederorr in result.failures:
                    self.stream.red("[  FAILED  ] %s" % failedtest)
                    self.stream.writeln()
            if errored:
                self.stream.red("[  ERRORED ] %d tests" % errored)
                for erroredtest, erorrmsg in result.errors:
                    self.stream.red("[  ERRORED ] %s" % erroredtest)
                    self.stream.writeln()

            self.stream.writeln()
            if failed:
                self.stream.write("%d ERRORED TEST" % failed)
            if errored:
                self.stream.write("%d ERRORED TEST" % errored)

        return result

TestResult类

class MyTestResult(unittest.TestResult):
    separator1 = '[----------] '
    separator2 = '[==========] '
    def __init__(self, stream=sys.stderr, descriptions=1, verbosity=1):
        unittest.TestResult.__init__(self)
        self.stream = stream
        self.showAll = verbosity > 1
        self.dots = verbosity == 1
        self.descriptions = descriptions

    def getDescription(self, test):
        if self.descriptions:
            return test.shortDescription() or str(test)
        else:
            return str(test)

    def startTest(self, test):
        self.stream.green('[ Run      ] ')
        self.stream.writeln(self.getDescription(test))
        unittest.TestResult.startTest(self, test)
        if self.showAll:
            self.stream.write(self.getDescription(test))
            self.stream.write(" ... ")

    def addSuccess(self, test):
        unittest.TestResult.addSuccess(self, test)
        if self.showAll:
            self.stream.writeln("ok")
        elif self.dots:
            self.stream.green('[       OK ] ')
            self.stream.writeln(self.getDescription(test))

    def addError(self, test, err):
        unittest.TestResult.addError(self, test, err)
        if self.showAll:
            self.stream.writeln("ERROR")
        elif self.dots:
            self.stream.write('E')

    def addFailure(self, test, err):
        unittest.TestResult.addFailure(self, test, err)
        if self.showAll:
            self.stream.writeln("FAIL")
        elif self.dots:
            self.stream.red('[  FAILED  ] ')
            self.stream.writeln(self.getDescription(test))
            self.stream.write(self._exc_info_to_string(err, test))