使用unittest为Python3脚本编写单元测试：模拟stdin求助

原Python脚本 #

你的目标脚本通过sys.stdin.read()读取标准输入并打印：

import sys

content = sys.stdin.read()
print(content)

问题 #

你需要为这段代码编写单元测试，模拟stdin输入，但之前的测试代码无法生效——运行测试时仍需要手动输入内容。

你尝试的测试代码 #

你之前写的测试代码如下：

import unittest
from unittest.mock import patch

import module_under_test


class MyTestCase(unittest.TestCase):

    def setUp(self):
        # raw_input is untouched before test
        assert module_under_test.raw_input is __builtins__.raw_input

    def test_using_with(self):
        input_data = "123"
        expected = int(input_data)

        with patch.object(module_under_test, "raw_input", create=True, 
                return_value=expected):
            # create=True is needed as raw_input is not in the globals of 
            # module_under_test, but actually found in __builtins__ .
            actual = module_under_test.function()

        self.assertEqual(expected, actual)

    @patch.object(module_under_test, "raw_input", create=True)
    def test_using_decorator(self, raw_input):
        raw_input.return_value = input_data = "123"
        expected = int(input_data)

        actual = module_under_test.function()

        self.assertEqual(expected, actual)

    def tearDown(self):
        # raw input is restored after test
        assert module_under_test.raw_input is __builtins__.raw_input

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

问题根源 #

你之前的测试方向错了：原脚本用的是sys.stdin.read()读取输入，而你的测试代码在模拟raw_input（这是Python2的输入函数，Python3对应input()），两者完全无关，所以模拟不会生效，测试时还是会读取真实的标准输入。

另外，原脚本在模块导入阶段就会执行sys.stdin.read()，导致测试还没开始模拟就触发了输入读取，这也是问题之一。

修正方案 #

第一步：修改原脚本，封装成函数 #

先把原脚本的逻辑封装成函数，避免模块导入时自动执行输入读取：

import sys

def read_and_print():
    content = sys.stdin.read()
    print(content)

if __name__ == "__main__":
    read_and_print()

第二步：编写正确的测试代码 #

使用unittest.mock.patch模拟sys.stdin，同时捕获sys.stdout来验证输出：

import unittest
from unittest.mock import patch
from io import StringIO
import module_under_test

class TestStdinRead(unittest.TestCase):
    def test_read_and_print(self):
        # 定义模拟的输入内容
        test_input = "测试输入内容\n第二行测试文本"
        # 预期输出：输入内容加上print自动添加的换行符
        expected_output = test_input + "\n"

        # 模拟sys.stdin为StringIO对象，提供预设输入
        with patch("sys.stdin", StringIO(test_input)):
            # 捕获sys.stdout的输出
            captured_output = StringIO()
            with patch("sys.stdout", captured_output):
                # 执行要测试的函数
                module_under_test.read_and_print()
                # 获取捕获到的输出内容
                actual_output = captured_output.getvalue()

        # 断言输出符合预期
        self.assertEqual(actual_output, expected_output)

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

代码说明 #

1. StringIO模拟输入输出：用StringIO创建内存中的字符串流，模拟标准输入和输出，避免依赖真实的控制台。
2. patch上下文管理器：通过patch替换sys.stdin和sys.stdout，测试完成后自动恢复原对象。
3. 函数封装：原脚本封装成函数后，只有主动调用时才会执行输入读取，确保测试可以控制执行时机。

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