utils.py
2.72 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
#!/usr/bin/env python3
"""Utils generally useful for unittests."""
import sys
import os
from os.path import dirname, join, abspath
from subprocess import check_output, PIPE, STDOUT, CalledProcessError
# Base dir of project, contains subdirs "tests" and "oletools" and README.md
PROJECT_ROOT = dirname(dirname(dirname(abspath(__file__))))
# Directory with test data, independent of current working directory
DATA_BASE_DIR = join(PROJECT_ROOT, 'tests', 'test-data')
# Directory with source code
SOURCE_BASE_DIR = join(PROJECT_ROOT, 'oletools')
def call_and_capture(module, args=None, accept_nonzero_exit=False,
exclude_stderr=False):
"""
Run module as script, capturing and returning output and return code.
This is the best way to capture a module's stdout and stderr; trying to
modify sys.stdout/sys.stderr to StringIO-Buffers frequently causes trouble.
Only drawback sofar: stdout and stderr are merged into one (which is
what users see on their shell as well). When testing for json-compatible
output you should `exclude_stderr` to `False` since logging ignores stderr,
so unforseen warnings (e.g. issued by pypy) would mess up your json.
:param str module: name of module to test, e.g. `olevba`
:param args: arguments for module's main function
:param bool fail_nonzero: Raise error if command returns non-0 return code
:param bool exclude_stderr: Exclude output to `sys.stderr` from output
(e.g. if parsing output through json)
:returns: ret_code, output
:rtype: int, str
"""
# create a PYTHONPATH environment var to prefer our current code
env = os.environ.copy()
try:
env['PYTHONPATH'] = SOURCE_BASE_DIR + os.pathsep + \
os.environ['PYTHONPATH']
except KeyError:
env['PYTHONPATH'] = SOURCE_BASE_DIR
# hack: in python2 output encoding (sys.stdout.encoding) was None
# although sys.getdefaultencoding() and sys.getfilesystemencoding were ok
# TODO: maybe can remove this once branch
# "encoding-for-non-unicode-environments" is merged
if 'PYTHONIOENCODING' not in env:
env['PYTHONIOENCODING'] = 'utf8'
# ensure args is a tuple
my_args = tuple(args) if args else ()
ret_code = -1
try:
output = check_output((sys.executable, '-m', module) + my_args,
universal_newlines=True, env=env,
stderr=PIPE if exclude_stderr else STDOUT)
ret_code = 0
except CalledProcessError as err:
if accept_nonzero_exit:
ret_code = err.returncode
output = err.output
else:
print(err.output)
raise
return output, ret_code