Closes issue 10979. unittest buffering now works with class and module setup and teardown

1 files changed, 17 insertions, 1 deletions

diff --git a/Lib/unittest/suite.py b/Lib/unittest/suite.py
index 77ce089..38bd6b8 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.
     """
@@ -133,6 +138,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:
@@ -142,6 +148,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
@@ -167,6 +175,7 @@ class TestSuite(BaseTestSuite):
             return
         setUpModule = getattr(module, 'setUpModule', None)
         if setUpModule is not None:
+            _call_if_exists(result, '_setupStdout')
             try:
                 setUpModule()
             except Exception as e:
@@ -175,6 +184,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)
@@ -198,6 +209,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:
@@ -205,6 +217,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)
@@ -220,6 +234,7 @@ class TestSuite(BaseTestSuite):
 
         tearDownClass = getattr(previousClass, 'tearDownClass', None)
         if tearDownClass is not None:
+            _call_if_exists(result, '_setupStdout')
             try:
                 tearDownClass()
             except Exception as e:
@@ -228,7 +243,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):