如何使用Python pytest/unittest测试函数控制台print输出内容

核心思路 #

通过临时重定向系统标准输出sys.stdout到内存缓冲区，执行完待测函数后读取缓冲区内容，即可验证打印内容是否符合预期。

实现步骤 #

方法1：手动重定向（无需额外封装） #

1. 导入依赖模块sys和StringIO
2. 测试执行前保存原始sys.stdout，替换为内存缓冲区实例
3. 执行待测函数
4. 读取缓冲区内容验证结果
5. 无论测试是否成功，都要恢复原始sys.stdout，避免影响其他测试逻辑

示例代码如下：

import sys
from StringIO import StringIO
import os
import xml.etree.ElementTree as ET
import unittest

# 待测函数
def load_parse_xml(data_file):
    if os.path.isfile(data_file):
        print "Parsing..."
        data_parse = ET.ElementTree(file=data_file)
        return data_parse.getroot()
    else:
        print data_file + " --> File not exist"

class TestXMLParse(unittest.TestCase):
    def test_load_parse_xml(self):
        # 保存原始标准输出
        original_stdout = sys.stdout
        captured_output = StringIO()
        # 替换标准输出为内存缓冲区
        sys.stdout = captured_output
        try:
            data_file = os.getcwd() + "/xml_tests/xmltest.xml"
            load_parse_xml(data_file)
            # 获取输出内容，strip()用于去除print默认添加的末尾换行符
            actual_output = captured_output.getvalue().strip()
            # 验证输出完全匹配
            self.assertEqual(actual_output, "Parsing...")
        finally:
            # 恢复原始标准输出
            sys.stdout = original_stdout

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

方法2：封装为上下文管理器（复用性更高） #

可以把重定向逻辑封装成通用的上下文管理器，后续其他测试用例也可以直接使用：

class capture_stdout(object):
    def __enter__(self):
        self.original_stdout = sys.stdout
        self.buffer = StringIO()
        sys.stdout = self.buffer
        return self
    def __exit__(self, exc_type, exc_val, exc_tb):
        sys.stdout = self.original_stdout
    def get_stripped_output(self):
        return self.buffer.getvalue().strip()

使用时代码更简洁：

def test_load_parse_xml(self):
    data_file = os.getcwd() + "/xml_tests/xmltest.xml"
    with capture_stdout() as output:
        load_parse_xml(data_file)
    self.assertEqual(output.get_stripped_output(), "Parsing...")

注意事项 #

• Python 2的print语句默认会在输出末尾添加换行符，如果你不想用strip()，也可以直接验证输出内容为"Parsing...\n"
• 不要省略sys.stdout的恢复逻辑，否则后续所有控制台打印都会被捕获到缓冲区，不会显示在终端

内容的提问来源于stack exchange，提问作者meuhfunk