Python: 让单元测试输出像GoogleTest一样

文章目录

- 1. 目的
- 2. 原版 unittest 的输出
- 3. 仿 GoogleTest 的输出效果
- 4. 实现原理浅析
- - 传入 testRunner 参数
  - testRunner 参数应该满足的条件
  - 颜色高亮： ASCII 转义字符的使用
  - 测试用例输出文本内容的格式调整：仿googletest
- 5. 完整实现代码
- 6. 完整调用代码

在这里插入图片描述

1. 目的

习惯了 C++ 单元测试框架 GoogleTest 的输出：成功则输出为绿色，失败则输出为红色，不同的case之间有布局上的分隔。使用 Python 虽然提高了开发效率，但为了保证程序质量，尤其是考虑到规模化的效果时，写单元测试仍然是不二法宝。对于 Python 自带的 unittest 单元测试模块而言，在配色上过于朴素（测试用例成功、失败都显示同一种颜色），考虑改进，而改为 GoogleTest 的样式则相对来说是一个明确的、接受度较高的目标。

2. 原版 unittest 的输出

首先给出两个测试用例，分别是成功的和失败的，然后查看 unittest 的输出效果:

import unittest

def inc(x):
    return x + 1

class JustTest(unittest.TestCase):
    def test_t1(self):
        assert inc(3) == 5
    def test_t2(self):
        assert inc(1) == 2

在这里插入图片描述
可以看到颜色很朴素。

3. 仿 GoogleTest 的输出效果

绿色表示成功，红色表示失败，很醒目了：
在这里插入图片描述
对应的单元测试写法也很简单，只需要改两行代码：

引入定制化 unittest 输出的模块
开启 unittest 时，传入定制化模块中的的类作为参数

import unittest
import mytest  ## ~ 新增这句

def inc(x):
    return x + 1

class JustTest(unittest.TestCase):
    def test_t1(self):
        assert inc(3) == 5
    def test_t2(self):
        assert inc(1) == 2

if __name__ == '__main__':
    #unittest.main()  ## ~ 这句是老的写法
    unittest.main(testRunner=mytest.MyTestRunner()) ## ~ 新增这句

4. 实现原理浅析

传入 testRunner 参数

unittest 模块的 main 函数接受 testRunner 这一参数, 而 unittest.main() 其实是 unittest.TestProgram 类的实例，传入的 testRunner 参数作为 TestProgram 类的构造函数的参数：
在这里插入图片描述

testRunner 参数应该满足的条件

满足的条件是：不传入 testRunner 参数时默认的 testRunner 对应的行为，也就是 runner.TextTestRunner 类。
在这里插入图片描述而 runner.TextTestRunner 类其实很简单，核心函数只有一个: run():

因此只要我们实现的 MyTestRunner 类能够兼容 TextTestRunner 即可。

颜色高亮： ASCII 转义字符的使用

在实现 log 的颜色打印时， ASCII 转义字符被使用到，网络资源很多可以自行查阅，效果大致如下：
在这里插入图片描述这里唯一要注意的是，要同时支持 Windows 和 Linux，MacOSX. Windows 下的颜色高亮实现方式略有不同，因而需要两份实现。

测试用例输出文本内容的格式调整：仿googletest

这一部分负责在每个测试用例运行的最开始，输出“套话”：[ Run ], 运行结果如果成功则输出 ok，不成功则输出 fail 或 error，就像 googletest 一样。


    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.red('[  ERRORED ] ')
            self.stream.writeln(self.getDescription(test))
            self.stream.write(self._exc_info_to_string(err, test))

    def addFailure(self, test, err):
        unittest.TestResult.addFailure(self, test, err)
        if self.showAll:
            self.stream.writeln("FAIL")
        elif self.dots:
            #self.stream.write(self._exc_info_to_string(err, test))
            content = self._exc_info_to_string(err, test)
            content_lines = content.split('\n')
            linenum_line = content_lines[1]
            file_desc, linenum_desc, testcase_desc = linenum_line.split(', ')
            linenum = linenum_desc.split(' ')[-1]
            filename = file_desc.split(' ')[-1][1:-1]
            result = '{:s}:{:s} Failure'.format(filename, linenum)
            self.stream.writeln(result)
            self.stream.writeln('Expected:')
            self.stream.writeln("\t" + content_lines[2].strip())
            self.stream.writeln('Actual:')
            self.stream.writeln("\t" + content_lines[3].strip())

            self.stream.red('[   FAILED ] ')
            self.stream.writeln(self.getDescription(test))

5. 完整实现代码

##################################################################################################################
# Make Python unittest output like googletest
# Author: ChrisZZ <imzhuo@foxmail.com>
# Homepage: <https://github.com/zchrissirhcz>
# --------------------
#
# Example usage:
# --------------------
# import unittest
# from mytest import MyTestRunner
#
# def inc(x):
#     return x + 1
#
# class JustTest(unittest.TestCase):
#     def test_answer(self):
#         assert inc(3) == 5
#
# if __name__ == '__main__':
#     unittest.main(testRunner=MyTestRunner())
#
# References
# --------------------
# https://www.cnblogs.com/coderzh/archive/2010/08/23/custom-python-unittestoutput-as-gtest.html
# https://github.com/xxhfg/youlook/blob/8d30e008260540902dacf5fbc5d4015bd68dc13e/libs/ColorUnittest/myunittest.py
# https://github.com/sndnyang/Tools/blob/d0bc2a7f5aa645048bd2fbaa28a43b1cccfac323/colortest.py
#
##################################################################################################################

import unittest
import time
import sys


import os

if os.name == 'nt':
    import ctypes

    ## {{{ 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.

    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

    class _ColorWritelnDecorator:
        """Used to decorate file-like objects with a handy 'writeln' method"""
        def __init__(self, stream):
            self.stream = stream
            if os.name == 'nt':
                self.std_out_handle = ctypes.windll.kernel32.GetStdHandle(STD_OUTPUT_HANDLE)

        def __getattr__(self, name):
            return getattr(self.stream, name)

        def yellow(self, msg):
            set_color(FOREGROUND_YELLOW | FOREGROUND_INTENSITY)
            self.write(msg)
            set_color(FOREGROUND_WHITE)

        def writeln(self, msg=None):
            if msg:
                self.write(msg)
            self.write('\n')

        def red(self, msg):
            set_color(FOREGROUND_RED | FOREGROUND_INTENSITY)
            self.write(msg)
            set_color(FOREGROUND_WHITE)

        def green(self, msg):
            set_color(FOREGROUND_GREEN | FOREGROUND_INTENSITY)
            self.write(msg)
            set_color(FOREGROUND_WHITE)

else:
    # https://en.wikipedia.org/wiki/ANSI_escape_code
    black = (1, 1, 1)
    red = (222, 56, 43)
    green = (57, 181, 74)
    yellow = (255, 199, 6)
    blue = (0, 111, 184)
    magenta = (118, 38, 113)
    cyan = (44, 181, 233)
    white = (204, 204, 204)
    bright_black = (128, 128, 128)
    bright_red = (255, 0, 0)
    bright_green = (0, 255, 0)
    bright_yellow = (255, 255, 0)
    bright_blue = (0, 0, 255)
    bright_magenta = (255, 0, 255)
    bright_cyan = (0, 255, 255)
    bright_white = (255, 255, 255)

    class _ColorWritelnDecorator:
        """Used to decorate file-like objects with a handy 'writeln' method"""
        def __init__(self, stream):
            self.stream = stream
            if os.name == 'nt':
                self.std_out_handle = ctypes.windll.kernel32.GetStdHandle(STD_OUTPUT_HANDLE)

        def __getattr__(self, name):
            return getattr(self.stream, name)

        def set_color(self, color):
            #stream.write("\033[38;2;{};{};{}m{} \033[38;2;255;255;255m".format(255, 0, 0, msg))
            r, g, b = color
            self.stream.write("\033[38;2;{};{};{}m ".format(r, g, b))

        def yellow(self, msg):
            self.set_color(yellow)
            self.write(msg)
            self.set_color(white)

        def writeln(self, msg=None):
            if msg:
                self.write(msg)
            self.write('\n')

        def red(self, msg):
            self.set_color(red)
            self.write(msg)
            self.set_color(white)

        def green(self, msg):
            self.set_color(green)
            self.write(msg)
            self.set_color(white)

import re

pattern = re.compile('File "(.+)",', re.IGNORECASE)

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.red('[  ERRORED ] ')
            self.stream.writeln(self.getDescription(test))
            self.stream.write(self._exc_info_to_string(err, test))

    def addFailure(self, test, err):
        unittest.TestResult.addFailure(self, test, err)
        if self.showAll:
            self.stream.writeln("FAIL")
        elif self.dots:
            #self.stream.write(self._exc_info_to_string(err, test))
            content = self._exc_info_to_string(err, test)
            content_lines = content.split('\n')
            linenum_line = content_lines[1]
            file_desc, linenum_desc, testcase_desc = linenum_line.split(', ')
            linenum = linenum_desc.split(' ')[-1]
            filename = file_desc.split(' ')[-1][1:-1]
            result = '{:s}:{:s} Failure'.format(filename, linenum)
            self.stream.writeln(result)
            self.stream.writeln('Expected:')
            self.stream.writeln("\t" + content_lines[2].strip())
            self.stream.writeln('Actual:')
            self.stream.writeln("\t" + content_lines[3].strip())

            self.stream.red('[   FAILED ] ')
            self.stream.writeln(self.getDescription(test))


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.green(result.separator2)
        self.stream.writeln('Your Unit Tests Start')

        startTime = time.time()
        test(result)

        stopTime = time.time()
        timeTaken = stopTime - startTime
        self.stream.green(result.separator2)
        run = result.testsRun
        self.stream.writeln("Run %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:
                    match = pattern.findall(failederorr)
                    if match:
                        src_file = match[0]
                    self.stream.red("[  FAILED  ] %s in %s" % (failedtest, src_file))
                    self.stream.writeln()
            if errored:
                self.stream.red("[  ERRORED ] %d tests, listed below:" % errored)
                self.stream.writeln()
                for erroredtest, erorrmsg in result.errors:
                    match = pattern.findall(erorrmsg)
                    if match:
                        src_file = match[0]
                    self.stream.red("[  ERRORED ] %s in %s" % (erroredtest, src_file))
                    self.stream.writeln()

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

        return result

6. 完整调用代码

import unittest
import mytest

def inc(x):
    return x + 1

class JustTest(unittest.TestCase):
    def test_t1(self):
        assert inc(3) == 5
    def test_t2(self):
        assert inc(1) == 2

if __name__ == '__main__':
    #unittest.main()
    unittest.main(testRunner=mytest.MyTestRunner())