使用unittest.mock无法捕获__main__.py中print函数的输出

核心问题 #

你的代码捕获不到输出，主要有两个关键原因：

• __main__.py的入口代码未执行：Python中if __name__ == '__main__':块里的代码只有在脚本被直接运行时才会执行，通过import导入模块时不会触发这块逻辑，所以你导入__main__模块后，里面的核心代码根本没跑，自然没有输出。
• 标准输出替换时机不当：手动替换sys.stdout的时机在导入模块之后，且如果__main__.py里的print已经提前绑定了原始sys.stdout，替换后也无法捕获输出。

修复方案 #

步骤1：重构__main__.py

先把__main__.py里的核心逻辑封装成独立函数（比如main()），让入口代码只负责调用这个函数：

# SIESTAstepper/__main__.py
import sys

def main():
    # 原来写在if __name__ == '__main__':里的所有逻辑
    print("测试输出")
    # 处理sys.argv的业务逻辑...

if __name__ == '__main__':
    main()

步骤2：修改测试函数

用unittest.mock.patch同时捕获sys.stdout和sys.argv，然后直接调用封装好的main()函数：

import io
from unittest.mock import patch
from SIESTAstepper.__main__ import main

def main_tester(command):
    fake_command()
    with patch('sys.stdout', new_callable=io.StringIO) as captured_output, \
         patch('sys.argv', command.split(" ")):
        main()
        return captured_output.getvalue()

为什么这样有效？ #

• 直接调用main()函数，确保核心逻辑被执行，不再依赖__name__ == '__main__'的判断。
• 用patch管理sys.stdout的替换，能确保main()函数里的print使用的是被替换后的输出流，避免手动替换的时机问题。
• 同时在同一个with块中处理sys.argv和sys.stdout的patch，保证上下文一致。

备选方案（不重构__main__.py的情况） #

如果暂时不想重构__main__.py，可以通过模拟直接运行脚本的方式执行，但这种方式不够优雅：

import os
import sys
import io
from unittest.mock import patch

def main_tester(command):
    fake_command()
    with patch('sys.argv', command.split(" ")), \
         patch('sys.stdout', new_callable=io.StringIO) as captured_output:
        # 重新执行__main__.py的代码，绕过__name__判断
        module_path = os.path.join(os.path.dirname(__file__), '../SIESTAstepper/__main__.py')
        with open(module_path, 'r') as f:
            code = compile(f.read(), module_path, 'exec')
            exec(code, {'__name__': '__main__'})
        return captured_output.getvalue()

不过更推荐第一种重构的方式，让代码结构更清晰，测试也更简单。

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