Python单元测试读取标准输入模块时第二个测试输出为空问题

核心问题是模块顶层代码只会在第一次导入时执行。第一次测试用例运行时，导入main模块，模块内的print(len(input()) + ...)会立即执行，此时mock的输入生效，输出正常。但第二个测试用例执行时，Python已将main模块缓存到sys.modules中，即便尝试卸载模块，只要测试代码中仍存在对main的引用，缓存就无法彻底清除，模块顶层代码不会再次执行，因此没有输出。

1. 重构main.py，将逻辑封装为函数（推荐） #

把原顶层代码放到函数里，避免依赖模块导入时机，测试时可主动调用：

# main.py
def check_total_length():
    total = len(input()) + len(input()) + len(input())
    print(total == 10)

if __name__ == "__main__":
    check_total_length()

对应的测试用例调整为：

import unittest
from unittest.mock import patch
import main

class TestMain(unittest.TestCase):
    def test_total_length_equal_10(self):
        with patch('builtins.input', side_effect=["a", "abcdefgh", "I"]):
            with patch('builtins.print') as mock_print:
                main.check_total_length()
                mock_print.assert_called_once_with(True)
    
    def test_total_length_not_equal_10(self):
        with patch('builtins.input', side_effect=["a", "b", "c"]):
            with patch('builtins.print') as mock_print:
                main.check_total_length()
                mock_print.assert_called_once_with(False)

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

2. 若必须保留原main.py结构 #

可在每个测试用例前强制重新加载模块：

import unittest
from unittest.mock import patch
import sys
import importlib

class TestMain(unittest.TestCase):
    def setUp(self):
        # 每次测试前移除模块缓存并重新导入
        if 'main' in sys.modules:
            del sys.modules['main']
        self.main_module = importlib.import_module('main')
    
    def test_total_length_equal_10(self):
        with patch('builtins.input', side_effect=["a", "abcdefgh", "I"]):
            with patch('builtins.print') as mock_print:
                importlib.reload(self.main_module)
                mock_print.assert_called_once_with(True)
    
    def test_total_length_not_equal_10(self):
        with patch('builtins.input', side_effect=["a", "b", "c"]):
            with patch('builtins.print') as mock_print:
                importlib.reload(self.main_module)
                mock_print.assert_called_once_with(False)

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

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