Use own buffering mechanism to capture stdout / stderr for XML reports

Author: ikedam

tests/testsuite.py

@@ -3,6 +3,8 @@
 
 """Executable module to test unittest-xml-reporting.
 """
+import contextlib
+import io
 import sys
 
 from xmlrunner.unittest import unittest
@@ -55,6 +57,22 @@ def test_doctest_example(self):
                       'name="twice"'.encode('utf8'), output)
 
 
+@contextlib.contextmanager
+def capture_stdout_stderr():
+    """
+    context manager to capture stdout and stderr
+    """
+    orig_stdout = sys.stdout
+    orig_stderr = sys.stderr
+    sys.stdout = StringIO()
+    sys.stderr = StringIO()
+    try:
+        yield (sys.stdout, sys.stderr)
+    finally:
+        sys.stdout = orig_stdout
+        sys.stderr = orig_stderr
+
+
 class XMLTestRunnerTestCase(unittest.TestCase):
     """
     XMLTestRunner test case.
@@ -102,6 +120,9 @@ def test_runner_buffer_output_fail(self):
             print('should be printed')
             self.fail('expected to fail')
 
+        def test_output(self):
+            print('test message')
+
         def test_non_ascii_runner_buffer_output_fail(self):
             print(u'Where is the café ?')
             self.fail(u'The café could not be found')
@@ -270,7 +291,25 @@ def test_xmlrunner_non_ascii_failures(self):
         runner = xmlrunner.XMLTestRunner(
             stream=self.stream, output=outdir, verbosity=self.verbosity,
             **self.runner_kwargs)
-        runner.run(suite)
+
+        # allow output non-ascii letters to stdout
+        orig_stdout = sys.stdout
+        if getattr(sys.stdout, 'buffer', None):
+            # Python3
+            sys.stdout = io.TextIOWrapper(sys.stdout.buffer, encoding='utf-8')
+        else:
+            # Python2
+            import codecs
+            sys.stdout = codecs.getwriter("utf-8")(sys.stdout)
+
+        try:
+            runner.run(suite)
+        finally:
+            if getattr(sys.stdout, 'buffer', None):
+                # Python3
+                # Not to be closed when TextIOWrapper is disposed.
+                sys.stdout.detach()
+            sys.stdout = orig_stdout
         outdir.seek(0)
         output = outdir.read()
         self.assertIn(
@@ -296,10 +335,52 @@ def test_xmlrunner_buffer_output_pass(self):
     def test_xmlrunner_buffer_output_fail(self):
         suite = unittest.TestSuite()
         suite.addTest(self.DummyTest('test_runner_buffer_output_fail'))
+        # --buffer option
+        self.runner_kwargs['buffer'] = True
         self._test_xmlrunner(suite)
         testsuite_output = self.stream.getvalue()
         self.assertIn('should be printed', testsuite_output)
 
+    def test_xmlrunner_output_without_buffer(self):
+        suite = unittest.TestSuite()
+        suite.addTest(self.DummyTest('test_output'))
+        with capture_stdout_stderr() as r:
+            self._test_xmlrunner(suite)
+        output_from_test = r[0].getvalue()
+        self.assertIn('test message', output_from_test)
+
+    def test_xmlrunner_output_with_buffer(self):
+        suite = unittest.TestSuite()
+        suite.addTest(self.DummyTest('test_output'))
+        # --buffer option
+        self.runner_kwargs['buffer'] = True
+        with capture_stdout_stderr() as r:
+            self._test_xmlrunner(suite)
+        output_from_test = r[0].getvalue()
+        self.assertNotIn('test message', output_from_test)
+
+    def test_xmlrunner_stdout_stderr_recovered_without_buffer(self):
+        orig_stdout = sys.stdout
+        orig_stderr = sys.stderr
+        suite = unittest.TestSuite()
+        suite.addTest(self.DummyTest('test_pass'))
+        self._test_xmlrunner(suite)
+        self.assertIs(orig_stdout, sys.stdout)
+        self.assertIs(orig_stderr, sys.stderr)
+
+    def test_xmlrunner_stdout_stderr_recovered_with_buffer(self):
+        orig_stdout = sys.stdout
+        orig_stderr = sys.stderr
+        suite = unittest.TestSuite()
+        suite.addTest(self.DummyTest('test_pass'))
+        # --buffer option
+        self.runner_kwargs['buffer'] = True
+        self._test_xmlrunner(suite)
+        self.assertIs(orig_stdout, sys.stdout)
+        self.assertIs(orig_stderr, sys.stderr)
+        suite = unittest.TestSuite()
+        suite.addTest(self.DummyTest('test_pass'))
+
     @unittest.skipIf(not hasattr(unittest.TestCase, 'subTest'),
                      'unittest.TestCase.subTest not present.')
     def test_unittest_subTest_fail(self):

xmlrunner/result.py

@@ -1,4 +1,5 @@
 
+import io
 import os
 import sys
 import datetime
@@ -82,6 +83,37 @@ def testcase_name(test_method):
     return result
 
 
+class _DuplicateWriter(io.RawIOBase):
+    """
+    Duplicate output from the first handle to the second handle
+
+    The second handle is expected to be a StringIO and not to block.
+    """
+
+    def __init__(self, first, second):
+        super(_DuplicateWriter, self).__init__()
+        self._first = first
+        self._second = second
+
+    def flush(self):
+        self._first.flush()
+        self._second.flush()
+
+    def writable(self):
+        return True
+
+    def writelines(self, lines):
+        self._first.writelines(lines)
+        self._second.writelines(lines)
+
+    def write(self, b):
+        len = self._first.write(b)
+
+        # expected to always succeed to write
+        self._second.write(b[:len])
+        return len
+
+
 class _TestInfo(object):
     """
     This class keeps useful information about the execution of a
@@ -152,9 +184,12 @@ class _XMLTestResult(_TextTestResult):
     def __init__(self, stream=sys.stderr, descriptions=1, verbosity=1,
                  elapsed_times=True, properties=None, infoclass=None):
         _TextTestResult.__init__(self, stream, descriptions, verbosity)
-        self.buffer = True  # we are capturing test output
         self._stdout_data = None
         self._stderr_data = None
+        self._stdout_capture = StringIO()
+        self.__stdout_saved = None
+        self._stderr_capture = StringIO()
+        self.__stderr_saved = None
         self.successes = []
         self.callback = None
         self.elapsed_times = elapsed_times
@@ -202,14 +237,35 @@ def startTest(self, test):
             self.stream.write('  ' + self.getDescription(test))
             self.stream.write(" ... ")
 
+    def _setupStdout(self):
+        """
+        Capture stdout / stderr by replacing sys.stdout / sys.stderr
+        """
+        super(_XMLTestResult, self)._setupStdout()
+        self.__stdout_saved = sys.stdout
+        sys.stdout = _DuplicateWriter(sys.stdout, self._stdout_capture)
+        self.__stderr_saved = sys.stderr
+        sys.stderr = _DuplicateWriter(sys.stderr, self._stderr_capture)
+
+    def _restoreStdout(self):
+        """
+        Stop capturing stdout / stderr and recover sys.stdout / sys.stderr
+        """
+        if self.__stdout_saved:
+            sys.stdout = self.__stdout_saved
+            self.__stdout_saved = None
+        if self.__stderr_saved:
+            sys.stderr = self.__stderr_saved
+            self.__stderr_saved = None
+        self._stdout_capture.seek(0)
+        self._stdout_capture.truncate()
+        self._stderr_capture.seek(0)
+        self._stderr_capture.truncate()
+        super(_XMLTestResult, self)._restoreStdout()
+
     def _save_output_data(self):
-        # Only try to get sys.stdout and sys.sterr as they not be
-        # StringIO yet, e.g. when test fails during __call__
-        try:
-            self._stdout_data = sys.stdout.getvalue()
-            self._stderr_data = sys.stderr.getvalue()
-        except AttributeError:
-            pass
+        self._stdout_data = self._stdout_capture.getvalue()
+        self._stderr_data = self._stderr_capture.getvalue()
 
     def stopTest(self, test):
         """
@@ -526,16 +582,8 @@ def _exc_info_to_string(self, err, test):
             msgLines = traceback.format_exception(exctype, value, tb)
 
         if self.buffer:
-            # Only try to get sys.stdout and sys.sterr as they not be
-            # StringIO yet, e.g. when test fails during __call__
-            try:
-                output = sys.stdout.getvalue()
-            except AttributeError:
-                output = None
-            try:
-                error = sys.stderr.getvalue()
-            except AttributeError:
-                error = None
+            output = self._stdout_capture.getvalue()
+            error = self._stdout_capture.getvalue()
             if output:
                 if not output.endswith('\n'):
                     output += '\n'

xmlrunner/runner.py

@@ -52,6 +52,7 @@ def run(self, test):
             # Prepare the test execution
             result = self._make_result()
             result.failfast = self.failfast
+            result.buffer = self.buffer
             if hasattr(test, 'properties'):
                 # junit testsuite properties
                 result.properties = test.properties