Python unittest：如何临时将标准输出消息重定向到缓冲区并测试其内容？

Question

I would like to capture messages sent to stdout (our stderr ) temporarily during a test and assert if some string patterns occurred in these messages:

import unittest
class SomeTest(unittest.TestCase):
    def test_stdout(self):
        output = ""
        function_that_writes_to_stdout()
        # How to capture stdout in output temporarily?
        self.assertIn("some message", output)

I found a similar question , but the accepted answer suggests to capture messages sent to stdout for all test cases.

Yes, I know that it's not very smart to unittest messages sent to stdout. And yes, I also know that it'd be better to use logging in combination with assertLogs . Let's assume that both options are not available at this stage.

Answer 1

Solution 1. The following worked for me:

import io
import unittest
from contextlib import redirect_stdout

class Test(unittest.TestCase):
    def test_stdout(self):
        buf = io.StringIO()
        with redirect_stdout(buf):
            print("foo!")
        self.assertIn("foo", buf.getvalue())

buf.getvalue() will contain the entire output, including \\n characters.

Solution 2. To mimic the behavior of assertLogs , one can extend unittest.TestCase by a method assertStdout as follows.

class StdoutRedirectionContext():
    class ListIO():
        def __init__(self):
            # Container for messages sent to stdout.
            self.output = []
        def write(self, s):
            # Filter empty strings or naked newline characters.
            if s in ("\n", ""): return
            self.output.append(s)

    def __enter__(self):
        self._buf = self.ListIO()
        self._ctx = redirect_stdout(self._buf)
        self._ctx.__enter__()
        return self._buf

    def __exit__(self, exc_type, exc_value, exc_traceback):
        self._ctx.__exit__(exc_type, exc_value, exc_traceback)

class TestCase(unittest.TestCase):
    def assertStdout(self):
        return StdoutRedirectionContext()

Here, StdoutRedirectionContext acts a context manager, and the single messages will be collected in the output list. The extended TestCase can be used as follows to assert messages on stdout:

class AnotherTest(TestCase):
    def test_stdout(self):
        with self.assertStdout() as cm:
            print("foo!")
            print("bar!")
        self.assertIn("foo!", cm.output)
        self.assertIn("baz!", cm.output)

The above yields the following output:

======================================================================
FAIL: test_stdout (__main__.AnotherTest)
----------------------------------------------------------------------
Traceback (most recent call last):
  File "executor_test.py", line 440, in test_stdout
    self.assertIn("baz!", cm.output)
AssertionError: 'baz!' not found in ['foo!', 'bar!']