如何在Python中使用nosetest/unittest断言输出？

Question

我正在为一个函数写测试，像下面这样：

def foo():
    print 'hello world!'

所以当我想测试这个函数时，代码会像这样：

import sys
from foomodule import foo
def test_foo():
    foo()
    output = sys.stdout.getline().strip() # because stdout is an StringIO instance
    assert output == 'hello world!'

但是如果我用 -s 参数运行 nosetests，测试就会崩溃。我该如何用 unittest 或 nose 模块来捕捉输出呢？

Answer 1

从2.7版本开始，你不再需要重新设置 sys.stdout 了，这个功能是通过 buffer 标志 提供的。而且，这也是nosetest的默认行为。

下面是一个在没有缓冲的情况下失败的示例：

import sys
import unittest

def foo():
    print 'hello world!'

class Case(unittest.TestCase):
    def test_foo(self):
        foo()
        if not hasattr(sys.stdout, "getvalue"):
            self.fail("need to run in buffered mode")
        output = sys.stdout.getvalue().strip() # because stdout is an StringIO instance
        self.assertEquals(output,'hello world!')

你可以通过 unit2 命令行标志 -b 或 --buffer 来设置缓冲，或者在 unittest.main 的选项中设置。想要关闭这个功能，可以使用 nosetest 的 --nocapture 标志。

if __name__=="__main__":   
    assert not hasattr(sys.stdout, "getvalue")
    unittest.main(module=__name__, buffer=True, exit=False)
    #.
    #----------------------------------------------------------------------
    #Ran 1 test in 0.000s
    #
    #OK
    assert not hasattr(sys.stdout, "getvalue")

    unittest.main(module=__name__, buffer=False)
    #hello world!
    #F
    #======================================================================
    #FAIL: test_foo (__main__.Case)
    #----------------------------------------------------------------------
    #Traceback (most recent call last):
    #  File "test_stdout.py", line 15, in test_foo
    #    self.fail("need to run in buffered mode")
    #AssertionError: need to run in buffered mode
    #
    #----------------------------------------------------------------------
    #Ran 1 test in 0.002s
    #
    #FAILED (failures=1)

Answer 2

如果你真的想这么做，可以在测试期间重新指定输出的地方。

def test_foo():
    import sys
    from foomodule import foo
    from StringIO import StringIO

    saved_stdout = sys.stdout
    try:
        out = StringIO()
        sys.stdout = out
        foo()
        output = out.getvalue().strip()
        assert output == 'hello world!'
    finally:
        sys.stdout = saved_stdout

不过如果是我写这段代码，我会更喜欢给foo函数加一个可选的out参数。

def foo(out=sys.stdout):
    out.write("hello, world!")

这样一来，测试就简单多了：

def test_foo():
    from foomodule import foo
    from StringIO import StringIO

    out = StringIO()
    foo(out=out)
    output = out.getvalue().strip()
    assert output == 'hello world!'

Answer 3

我使用这个上下文管理器来捕获输出。它其实和其他一些答案用的技术是一样的，都是暂时替换sys.stdout。我更喜欢上下文管理器，因为它把所有的管理工作都封装在一个函数里，这样我就不需要重新写任何的try-finally代码，也不需要为了这个特意写设置和清理的函数。

import sys
from contextlib import contextmanager
from StringIO import StringIO

@contextmanager
def captured_output():
    new_out, new_err = StringIO(), StringIO()
    old_out, old_err = sys.stdout, sys.stderr
    try:
        sys.stdout, sys.stderr = new_out, new_err
        yield sys.stdout, sys.stderr
    finally:
        sys.stdout, sys.stderr = old_out, old_err

用法如下：

with captured_output() as (out, err):
    foo()
# This can go inside or outside the `with` block
output = out.getvalue().strip()
self.assertEqual(output, 'hello world!')

而且，由于在退出with块时会恢复原来的输出状态，我们可以在同一个函数里设置第二个捕获块，这在使用设置和清理函数时是做不到的，手动写try-finally块时也会显得很繁琐。当测试的目标是比较两个函数的结果而不是和某个预先计算的值比较时，这种能力就特别有用了。