unittest: create test util dir; move output capture there

Want to use the OutputCapture in another unit test. There will probably be more such test helpers over time that need to be imported from different unit tests.

unittest: create test util dir; move output capture there
Want to use the OutputCapture in another unit test. There will probably be more such test helpers over time that need to be imported from different unit tests.
Christian Herdtweck
1 parent 1e8e5737
Showing 3 changed files with 54 additions and 32 deletions
tests/msodde_doc/test_basic.py
tests/test_utils/__init__.py
tests/test_utils/output_capture.py
@@ -10,16 +10,10 @@ from __future__ import print_function
 import unittest
 from oletools import msodde
+from tests.test_utils import OutputCapture
 import shlex
 from os.path import join, dirname, normpath
-import sys
-
-# python 2/3 version conflict:
-if sys.version_info.major <= 2:
-    from StringIO import StringIO
-    #from io import BytesIO as StringIO - try if print() gives UnicodeError
-else:
-    from io import StringIO
+from traceback import print_exc
 # base directory for test input
@@ -74,6 +68,7 @@ class TestReturnCode(unittest.TestCase):
             return_code = msodde.main(args)
         except Exception:
             have_exception = True
+            print_exc()
         except SystemExit as se:     # sys.exit() was called
             return_code = se.code
             if se.code is None:
@@ -85,30 +80,6 @@ class TestReturnCode(unittest.TestCase):
                                      return_code))
-class OutputCapture:
-    """ context manager that captures stdout """
-
-    def __init__(self):
-        self.output = StringIO()   # in py2, this actually is BytesIO
-
-    def __enter__(self):
-        sys.stdout = self.output
-        return self
-
-    def __exit__(self, exc_type, exc_value, traceback):
-        sys.stdout = sys.__stdout__    # re-set to original
-
-        if exc_type:   # there has been an error
-            print('Got error during output capture!')
-            print('Print captured output and re-raise:')
-            for line in self.output.getvalue().splitlines():
-                print(line.rstrip())  # print output before re-raising
-
-    def __iter__(self):
-        for line in self.output.getvalue().splitlines():
-            yield line.rstrip()   # remove newline at end of line
-
-
 class TestDdeInDoc(unittest.TestCase):
     def get_dde_from_output(self, capturer):
+from .output_capture import OutputCapture
+""" class OutputCapture to test what scripts print to stdout """
+
+from __future__ import print_function
+import sys
+
+
+# python 2/3 version conflict:
+if sys.version_info.major <= 2:
+    from StringIO import StringIO
+else:
+    from io import StringIO
+
+class OutputCapture:
+    """ context manager that captures stdout
+
+    use as follows::
+
+        with OutputCapture() as capturer:
+            run_my_script(some_args)
+
+        # either test line-by-line ...
+        for line in capturer:
+            some_test(line)
+        # ...or test all output in one go
+        some_test(capturer.buffer.getvalue())
+
+    """
+
+    def __init__(self):
+        self.buffer = StringIO()
+        self.orig_stdout = None
+
+    def __enter__(self):
+        # replace sys.stdout with own buffer.
+        self.orig_stdout = sys.stdout
+        sys.stdout = self.buffer
+        return self
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        sys.stdout = self.orig_stdout    # re-set to original
+
+        if exc_type:   # there has been an error
+            print('Got error during output capture!')
+            print('Print captured output and re-raise:')
+            for line in self.buffer.getvalue().splitlines():
+                print(line.rstrip())  # print output before re-raising
+
+    def __iter__(self):
+        for line in self.buffer.getvalue().splitlines():
+            yield line.rstrip()   # remove newline at end of line