bpo-13236: Flush the output stream more often in unittest (GH-29864)

It can prevent some losses when output to buffered stream.

2 files changed, 46 insertions, 5 deletions

diff --git a/Lib/unittest/test/test_program.py b/Lib/unittest/test/test_program.py
index 93a6109..cc02b22 100644
--- a/Lib/unittest/test/test_program.py
+++ b/Lib/unittest/test/test_program.py
@@ -6,6 +6,7 @@ import subprocess
 from test import support
 import unittest
 import unittest.test
+from .test_result import BufferedWriter
 
 
 class Test_TestProgram(unittest.TestCase):
@@ -104,30 +105,39 @@ class Test_TestProgram(unittest.TestCase):
                           program.testNames)
 
     def test_NonExit(self):
+        stream = BufferedWriter()
         program = unittest.main(exit=False,
                                 argv=["foobar"],
-                                testRunner=unittest.TextTestRunner(stream=io.StringIO()),
+                                testRunner=unittest.TextTestRunner(stream=stream),
                                 testLoader=self.FooBarLoader())
         self.assertTrue(hasattr(program, 'result'))
+        self.assertIn('\nFAIL: testFail ', stream.getvalue())
+        self.assertTrue(stream.getvalue().endswith('\n\nFAILED (failures=1)\n'))
 
 
     def test_Exit(self):
+        stream = BufferedWriter()
         self.assertRaises(
             SystemExit,
             unittest.main,
             argv=["foobar"],
-            testRunner=unittest.TextTestRunner(stream=io.StringIO()),
+            testRunner=unittest.TextTestRunner(stream=stream),
             exit=True,
             testLoader=self.FooBarLoader())
+        self.assertIn('\nFAIL: testFail ', stream.getvalue())
+        self.assertTrue(stream.getvalue().endswith('\n\nFAILED (failures=1)\n'))
 
 
     def test_ExitAsDefault(self):
+        stream = BufferedWriter()
         self.assertRaises(
             SystemExit,
             unittest.main,
             argv=["foobar"],
-            testRunner=unittest.TextTestRunner(stream=io.StringIO()),
+            testRunner=unittest.TextTestRunner(stream=stream),
             testLoader=self.FooBarLoader())
+        self.assertIn('\nFAIL: testFail ', stream.getvalue())
+        self.assertTrue(stream.getvalue().endswith('\n\nFAILED (failures=1)\n'))
 
 
 class InitialisableProgram(unittest.TestProgram):
diff --git a/Lib/unittest/test/test_result.py b/Lib/unittest/test/test_result.py
index c4a49b4..224a784 100644
--- a/Lib/unittest/test/test_result.py
+++ b/Lib/unittest/test/test_result.py
@@ -33,6 +33,22 @@ def bad_cleanup2():
     raise ValueError('bad cleanup2')
 
 
+class BufferedWriter:
+    def __init__(self):
+        self.result = ''
+        self.buffer = ''
+
+    def write(self, arg):
+        self.buffer += arg
+
+    def flush(self):
+        self.result += self.buffer
+        self.buffer = ''
+
+    def getvalue(self):
+        return self.result
+
+
 class Test_TestResult(unittest.TestCase):
     # Note: there are not separate tests for TestResult.wasSuccessful(),
     # TestResult.errors, TestResult.failures, TestResult.testsRun or
@@ -335,10 +351,13 @@ class Test_TestResult(unittest.TestCase):
         self.assertTrue(result.shouldStop)
 
     def testFailFastSetByRunner(self):
-        runner = unittest.TextTestRunner(stream=io.StringIO(), failfast=True)
+        stream = BufferedWriter()
+        runner = unittest.TextTestRunner(stream=stream, failfast=True)
         def test(result):
             self.assertTrue(result.failfast)
         result = runner.run(test)
+        stream.flush()
+        self.assertTrue(stream.getvalue().endswith('\n\nOK\n'))
 
 
 class Test_TextTestResult(unittest.TestCase):
@@ -462,6 +481,12 @@ class Test_TextTestResult(unittest.TestCase):
             self.fail('fail')
         def testError(self):
             raise Exception('error')
+        @unittest.expectedFailure
+        def testExpectedFailure(self):
+            self.fail('fail')
+        @unittest.expectedFailure
+        def testUnexpectedSuccess(self):
+            pass
         def testSubTestSuccess(self):
             with self.subTest('one', a=1):
                 pass
@@ -483,7 +508,7 @@ class Test_TextTestResult(unittest.TestCase):
                 raise self.tearDownError
 
     def _run_test(self, test_name, verbosity, tearDownError=None):
-        stream = io.StringIO()
+        stream = BufferedWriter()
         stream = unittest.runner._WritelnDecorator(stream)
         result = unittest.TextTestResult(stream, True, verbosity)
         test = self.Test(test_name)
@@ -496,6 +521,8 @@ class Test_TextTestResult(unittest.TestCase):
         self.assertEqual(self._run_test('testSkip', 1), 's')
         self.assertEqual(self._run_test('testFail', 1), 'F')
         self.assertEqual(self._run_test('testError', 1), 'E')
+        self.assertEqual(self._run_test('testExpectedFailure', 1), 'x')
+        self.assertEqual(self._run_test('testUnexpectedSuccess', 1), 'u')
 
     def testLongOutput(self):
         classname = f'{__name__}.{self.Test.__qualname__}'
@@ -507,6 +534,10 @@ class Test_TextTestResult(unittest.TestCase):
                          f'testFail ({classname}) ... FAIL\n')
         self.assertEqual(self._run_test('testError', 2),
                          f'testError ({classname}) ... ERROR\n')
+        self.assertEqual(self._run_test('testExpectedFailure', 2),
+                         f'testExpectedFailure ({classname}) ... expected failure\n')
+        self.assertEqual(self._run_test('testUnexpectedSuccess', 2),
+                         f'testUnexpectedSuccess ({classname}) ... unexpected success\n')
 
     def testDotsOutputSubTestSuccess(self):
         self.assertEqual(self._run_test('testSubTestSuccess', 1), '.')