fix unit tests

Samir Aguiar · Christian Herdtweck
1 parent aa24ab97
Showing 4 changed files with 9 additions and 7 deletions
oletools/olevba.py
oletools/rtfobj.py
tests/json/test_output.py
tests/rtfobj/test_is_rtf.py
@@ -293,6 +293,7 @@ from oletools.thirdparty.pyparsing.pyparsing import \
         infixNotation, ParserElement
 from oletools import ppt_parser
 from oletools import oleform
+from oletools import rtfobj
 # monkeypatch email to fix issue #32:
@@ -701,7 +701,7 @@ def is_rtf(arg, treat_str_as_data=False):
     if isinstance(arg, str):      # could be bytes, but we assume file name
         if treat_str_as_data:
             try:
-                return arg[:magic_len].encode('ascii', error='strict').lower()\
+                return arg[:magic_len].encode('ascii', errors='strict').lower()\
                     == RTF_MAGIC
             except UnicodeError:
                 return False
@@ -28,7 +28,7 @@ class TestValidJson(unittest.TestCase):
             for filename in filenames:
                 yield join(dirpath, filename)
-    def run_and_parse(self, program, args, print_output=False):
+    def run_and_parse(self, program, args, print_output=False, check_return_code=True):
         """ run single program with single file and parse output """
         return_code = None
         with OutputCapture() as capturer:       # capture stdout
@@ -38,7 +38,7 @@ class TestValidJson(unittest.TestCase):
                 return_code = 1   # would result in non-zero exit
             except SystemExit as se:
                 return_code = se.code or 0   # se.code can be None
-        if return_code is not 0:
+        if check_return_code and return_code is not 0:
             if print_output:
                 print('Command failed ({0}) -- not parsing output'
                       .format(return_code))
@@ -82,7 +82,8 @@ class TestValidJson(unittest.TestCase):
     def test_olevba_recurse(self):
         """ Test olevba.py with -r """
         json_data = self.run_and_parse(olevba.main,
-                                       ['-j', '-r', join(DATA_BASE_DIR, '*')])
+                                       ['-j', '-r', join(DATA_BASE_DIR, '*')],
+                                       check_return_code=False)
         self.assertNotEqual(len(json_data), 0,
                             msg='olevba[3] returned non-zero or no output')
         self.assertNotEqual(json_data[-1]['n_processed'], 0,
@@ -23,9 +23,9 @@ class TestIsRtf(unittest.TestCase):
     def test_bytes(self):
         """ test that is_rtf works with bytearray """
-        self.assertTrue(is_rtf(RTF_MAGIC + b'asasdffdfasdfasdfasdfasdf'), True)
-        self.assertTrue(is_rtf(RTF_MAGIC.upper() + b'asdffasdfasdasdff'), True)
-        self.assertFalse(is_rtf(b'asdfasdfasdfasdfasdfasdasdfffsdfsdfa'), True)
+        self.assertTrue(is_rtf(RTF_MAGIC + b'asasdffdfasdfasdfasdfasdf', True), True)
+        self.assertTrue(is_rtf(RTF_MAGIC.upper() + b'asdffasdfasdasdff', True), True)
+        self.assertFalse(is_rtf(b'asdfasdfasdfasdfasdfasdasdfffsdfsdfa', True), True)
     def test_tuple(self):
         """ test that is_rtf works with byte tuples """