test: Use call_and_capture in olevba tests

Running main() within same interpreter capturing SystemExit and sys.stdout/err always causes trouble (at least in unittest). Fork another python process and capture from there, that is easier and cleaner in my view.

test: Use call_and_capture in olevba tests
Running main() within same interpreter capturing SystemExit and sys.stdout/err always causes trouble (at least in unittest). Fork another python process and capture from there, that is easier and cleaner in my view.
Christian Herdtweck
1 parent b22b36c5
Showing 1 changed file with 39 additions and 87 deletions
tests/olevba/test_basic.py
@@ -3,44 +3,12 @@ Test basic functionality of olevba[3]
 """
 import unittest
-import sys
 import os
 from os.path import join
-from contextlib import contextmanager
-try:
-    from cStringIO import StringIO
-except ImportError:   # py3:
-    from io import StringIO
-if sys.version_info.major <= 2:
-    from oletools import olevba
-else:
-    from oletools import olevba3 as olevba
+import re
 # Directory with test data, independent of current working directory
-from tests.test_utils import DATA_BASE_DIR
-
-
-@contextmanager
-def capture_output():
-    """
-    Temporarily replace stdout/stderr with buffers to capture output.
-
-    Once we only support python>=3.4: this is already built into python as
-    :py:func:`contextlib.redirect_stdout`.
-
-    Not quite sure why, but seems to only work once per test function ...
-    """
-    orig_stdout = sys.stdout
-    orig_stderr = sys.stderr
-
-    try:
-        sys.stdout = StringIO()
-        sys.stderr = StringIO()
-        yield sys.stdout, sys.stderr
-
-    finally:
-        sys.stdout = orig_stdout
-        sys.stderr = orig_stderr
+from tests.test_utils import DATA_BASE_DIR, call_and_capture
 class TestOlevbaBasic(unittest.TestCase):
@@ -57,62 +25,48 @@ class TestOlevbaBasic(unittest.TestCase):
     def do_test_behaviour(self, filename):
         """Helper for test_{text,empty}_behaviour."""
         input_file = join(DATA_BASE_DIR, 'basic', filename)
-        ret_code = -1
-
-        # run olevba, capturing its output and return code
-        with capture_output() as (stdout, stderr):
-            with self.assertRaises(SystemExit) as raise_context:
-                olevba.main([input_file, ])
-            ret_code = raise_context.exception.code
-
-        # check that return code is 0
-        self.assertEqual(ret_code, 0)
-
-        # check there are only warnings in stderr
-        stderr = stderr.getvalue()
-        skip_line = False
-        for line in stderr.splitlines():
-            if skip_line:
-                skip_line = False
-                continue
-            self.assertTrue(line.startswith('WARNING ') or
-                            'ResourceWarning' in line,
-                            msg='Line "{}" in stderr is unexpected for {}'\
-                                .format(line.rstrip(), filename))
-            if 'ResourceWarning' in line:
-                skip_line = True
-        self.assertIn('not encrypted', stderr)
-
-        # check stdout
-        stdout = stdout.getvalue().lower()
-        self.assertIn(input_file.lower(), stdout)
-        self.assertIn('type: text', stdout)
-        self.assertIn('no suspicious', stdout)
-        self.assertNotIn('error', stdout)
-        self.assertNotIn('warn', stdout)
+        output, _ = call_and_capture('olevba', args=(input_file, ))
+
+        # check output
+        self.assertTrue(re.search(r'^Type:\s+Text\s*$', output, re.MULTILINE),
+                        msg='"Type: Text" not found in output:\n' + output)
+        self.assertTrue(re.search(r'^No suspicious .+ found.$', output,
+                                  re.MULTILINE),
+                        msg='"No suspicous...found" not found in output:\n' + \
+                            output)
+        self.assertNotIn('error', output.lower())
+
+        # check warnings
+        for line in output.splitlines():
+            if line.startswith('WARNING ') and 'encrypted' in line:
+                continue   # encryption warnings are ok
+            elif 'warn' in line.lower():
+                raise self.fail('Found "warn" in output line: "{}"'
+                                .format(line.rstrip()))
+        self.assertIn('not encrypted', output)
     def test_rtf_behaviour(self):
         """Test behaviour of olevba when presented with an rtf file."""
         input_file = join(DATA_BASE_DIR, 'msodde', 'RTF-Spec-1.7.rtf')
-        ret_code = -1
-
-        # run olevba, capturing its output and return code
-        with capture_output() as (stdout, stderr):
-            with self.assertRaises(SystemExit) as raise_context:
-                olevba.main([input_file, ])
-            ret_code = raise_context.exception.code
+        output, ret_code = call_and_capture('olevba', args=(input_file, ),
+                                            accept_nonzero_exit=True)
         # check that return code is olevba.RETURN_OPEN_ERROR
         self.assertEqual(ret_code, 5)
-        stdout = stdout.getvalue().lower()
-        self.assertNotIn('error', stdout)
-        self.assertNotIn('warn', stdout)
-        stderr = stderr.getvalue().lower()
-        self.assertIn('fileopenerror', stderr)
-        self.assertIn('is rtf', stderr)
-        self.assertIn('rtfobj.py', stderr)
-        self.assertIn('not encrypted', stderr)
+        # check output:
+        self.assertIn('FileOpenError', output)
+        self.assertIn('is RTF', output)
+        self.assertIn('rtfobj.py', output)
+        self.assertIn('not encrypted', output)
+
+        # check warnings
+        for line in output.splitlines():
+            if line.startswith('WARNING ') and 'encrypted' in line:
+                continue   # encryption warnings are ok
+            elif 'warn' in line.lower():
+                raise self.fail('Found "warn" in output line: "{}"'
+                                .format(line.rstrip()))
     def test_crypt_return(self):
         """
@@ -136,11 +90,9 @@ class TestOlevbaBasic(unittest.TestCase):
                 continue
             full_name = join(CRYPT_DIR, filename)
             for args in ADD_ARGS:
-                try:
-                    olevba.main(args + [full_name, ])
-                    self.fail('Olevba should have exited')
-                except SystemExit as sys_exit:
-                    ret_code = sys_exit.code or 0   # sys_exit.code can be None
+                _, ret_code = call_and_capture('olevba',
+                                               args=[full_name, ] + args,
+                                               accept_nonzero_exit=True)
                 self.assertEqual(ret_code, CRYPT_RETURN_CODE,
                                  msg='Wrong return code {} for args {}'\
                                      .format(ret_code, args + [filename, ]))