Issue #10979. unittest stdout buffering now works for class and module fixtures.

3 files changed, 91 insertions, 2 deletions

diff --git a/Lib/unittest/result.py b/Lib/unittest/result.py
index 1f66f1b..2cc17d7 100644
--- a/Lib/unittest/result.py
+++ b/Lib/unittest/result.py
@@ -60,6 +60,9 @@ class TestResult(object):
         "Called when the given test is about to be run"
         self.testsRun += 1
         self._mirrorOutput = False
+        self._setupStdout()
+
+    def _setupStdout(self):
         if self.buffer:
             if self._stderr_buffer is None:
                 self._stderr_buffer = StringIO()
@@ -75,6 +78,10 @@ class TestResult(object):
 
     def stopTest(self, test):
         """Called when the given test has been run"""
+        self._restoreStdout()
+        self._mirrorOutput = False
+
+    def _restoreStdout(self):
         if self.buffer:
             if self._mirrorOutput:
                 output = sys.stdout.getvalue()
@@ -94,7 +101,6 @@ class TestResult(object):
             self._stdout_buffer.truncate()
             self._stderr_buffer.seek(0)
             self._stderr_buffer.truncate()
-        self._mirrorOutput = False
 
     def stopTestRun(self):
         """Called once after all tests are executed.
diff --git a/Lib/unittest/suite.py b/Lib/unittest/suite.py
index c0ae2ae..633af5c 100644
--- a/Lib/unittest/suite.py
+++ b/Lib/unittest/suite.py
@@ -8,6 +8,11 @@ from . import util
 __unittest = True
 
 
+def _call_if_exists(parent, attr):
+    func = getattr(parent, attr, lambda: None)
+    func()
+
+
 class BaseTestSuite(object):
     """A simple test suite that doesn't provide class or module shared fixtures.
     """
@@ -136,6 +141,7 @@ class TestSuite(BaseTestSuite):
 
         setUpClass = getattr(currentClass, 'setUpClass', None)
         if setUpClass is not None:
+            _call_if_exists(result, '_setupStdout')
             try:
                 setUpClass()
             except Exception as e:
@@ -145,7 +151,8 @@ class TestSuite(BaseTestSuite):
                 className = util.strclass(currentClass)
                 errorName = 'setUpClass (%s)' % className
                 self._addClassOrModuleLevelException(result, e, errorName)
-
+            finally:
+                _call_if_exists(result, '_restoreStdout')
 
     def _get_previous_module(self, result):
         previousModule = None
@@ -170,6 +177,7 @@ class TestSuite(BaseTestSuite):
             return
         setUpModule = getattr(module, 'setUpModule', None)
         if setUpModule is not None:
+            _call_if_exists(result, '_setupStdout')
             try:
                 setUpModule()
             except Exception, e:
@@ -178,6 +186,8 @@ class TestSuite(BaseTestSuite):
                 result._moduleSetUpFailed = True
                 errorName = 'setUpModule (%s)' % currentModule
                 self._addClassOrModuleLevelException(result, e, errorName)
+            finally:
+                _call_if_exists(result, '_restoreStdout')
 
     def _addClassOrModuleLevelException(self, result, exception, errorName):
         error = _ErrorHolder(errorName)
@@ -201,6 +211,7 @@ class TestSuite(BaseTestSuite):
 
         tearDownModule = getattr(module, 'tearDownModule', None)
         if tearDownModule is not None:
+            _call_if_exists(result, '_setupStdout')
             try:
                 tearDownModule()
             except Exception as e:
@@ -208,6 +219,8 @@ class TestSuite(BaseTestSuite):
                     raise
                 errorName = 'tearDownModule (%s)' % previousModule
                 self._addClassOrModuleLevelException(result, e, errorName)
+            finally:
+                _call_if_exists(result, '_restoreStdout')
 
     def _tearDownPreviousClass(self, test, result):
         previousClass = getattr(result, '_previousTestClass', None)
@@ -223,6 +236,7 @@ class TestSuite(BaseTestSuite):
 
         tearDownClass = getattr(previousClass, 'tearDownClass', None)
         if tearDownClass is not None:
+            _call_if_exists(result, '_setupStdout')
             try:
                 tearDownClass()
             except Exception, e:
@@ -231,6 +245,8 @@ class TestSuite(BaseTestSuite):
                 className = util.strclass(previousClass)
                 errorName = 'tearDownClass (%s)' % className
                 self._addClassOrModuleLevelException(result, e, errorName)
+            finally:
+                _call_if_exists(result, '_restoreStdout')
 
 
 class _ErrorHolder(object):
diff --git a/Lib/unittest/test/test_result.py b/Lib/unittest/test/test_result.py
index ec3fadd..f0f3a63 100644
--- a/Lib/unittest/test/test_result.py
+++ b/Lib/unittest/test/test_result.py
@@ -496,5 +496,72 @@ class TestOutputBuffering(unittest.TestCase):
             self.assertEqual(result._original_stderr.getvalue(), expectedErrMessage)
             self.assertMultiLineEqual(message, expectedFullMessage)
 
+    def testBufferSetupClass(self):
+        result = unittest.TestResult()
+        result.buffer = True
+
+        class Foo(unittest.TestCase):
+            @classmethod
+            def setUpClass(cls):
+                1/0
+            def test_foo(self):
+                pass
+        suite = unittest.TestSuite([Foo('test_foo')])
+        suite(result)
+        self.assertEqual(len(result.errors), 1)
+
+    def testBufferTearDownClass(self):
+        result = unittest.TestResult()
+        result.buffer = True
+
+        class Foo(unittest.TestCase):
+            @classmethod
+            def tearDownClass(cls):
+                1/0
+            def test_foo(self):
+                pass
+        suite = unittest.TestSuite([Foo('test_foo')])
+        suite(result)
+        self.assertEqual(len(result.errors), 1)
+
+    def testBufferSetUpModule(self):
+        result = unittest.TestResult()
+        result.buffer = True
+
+        class Foo(unittest.TestCase):
+            def test_foo(self):
+                pass
+        class Module(object):
+            @staticmethod
+            def setUpModule():
+                1/0
+
+        Foo.__module__ = 'Module'
+        sys.modules['Module'] = Module
+        self.addCleanup(sys.modules.pop, 'Module')
+        suite = unittest.TestSuite([Foo('test_foo')])
+        suite(result)
+        self.assertEqual(len(result.errors), 1)
+
+    def testBufferTearDownModule(self):
+        result = unittest.TestResult()
+        result.buffer = True
+
+        class Foo(unittest.TestCase):
+            def test_foo(self):
+                pass
+        class Module(object):
+            @staticmethod
+            def tearDownModule():
+                1/0
+
+        Foo.__module__ = 'Module'
+        sys.modules['Module'] = Module
+        self.addCleanup(sys.modules.pop, 'Module')
+        suite = unittest.TestSuite([Foo('test_foo')])
+        suite(result)
+        self.assertEqual(len(result.errors), 1)
+
+
 if __name__ == '__main__':
     unittest.main()