diff --git a/tests/common/test_encoding_handler.py b/tests/common/test_encoding_handler.py index f4bdb90..1e074e7 100644 --- a/tests/common/test_encoding_handler.py +++ b/tests/common/test_encoding_handler.py @@ -1,47 +1,70 @@ """Test common.ensure_stdout_handles_unicode""" +from __future__ import print_function + import unittest import sys -from subprocess import check_call +from subprocess import check_call, CalledProcessError from tempfile import mkstemp import os from os.path import isfile +from contextlib import contextmanager + +FILE_TEXT = u'The unicode check mark is \u2713.\n' + +@contextmanager +def temp_file(just_name=True): + """Context manager that creates temp file and deletes it in the end""" + tmp_descriptor = None + tmp_name = None + tmp_handle = None + try: + tmp_descriptor, tmp_name = mkstemp() + + # we create our own file handle since we want to be able to close the + # file and open it again for reading. + # We keep the os-level descriptor open so file name is still reserved + # for us + if just_name: + yield tmp_name + else: + tmp_handle = open(tmp_name, 'wb') + yield tmp_handle, tmp_name + except Exception: + raise + finally: + if tmp_descriptor is not None: + os.close(tmp_descriptor) + if tmp_handle is not None: + tmp_handle.close() + if tmp_name is not None and isfile(tmp_name): + os.unlink(tmp_name) class TestEncodingHandler(unittest.TestCase): """Tests replacing stdout encoding in various scenarios""" - def test_base(self): + def test_print(self): """Test regular unicode output not raise error""" check_call('{python} {this_file} print'.format(python=sys.executable, this_file=__file__), shell=True) - def test_redirect(self): + def test_print_redirect(self): """ Test redirection of unicode output to files does not raise error TODO: test this on non-linux OSs """ - tmp_handle = None - tmp_name = None - try: - tmp_handle, tmp_name = mkstemp() + with temp_file() as tmp_file: check_call('{python} {this_file} print > {tmp_file}' .format(python=sys.executable, this_file=__file__, tmp_file=tmp_file), shell=True) - except Exception: - raise - finally: - if tmp_handle is not None: - os.close(tmp_handle) - if tmp_name is not None and isfile(tmp_name): - os.unlink(tmp_name) @unittest.skipIf(not sys.platform.startswith('linux'), 'Only tested on linux sofar') - def test_no_lang(self): + def test_print_no_lang(self): """ Test redirection of unicode output to files does not raise error @@ -51,28 +74,133 @@ class TestEncodingHandler(unittest.TestCase): .format(python=sys.executable, this_file=__file__), shell=True) + def test_uopen(self): + """Test that uopen in a nice environment is ok""" + with temp_file(False) as (tmp_handle, tmp_file): + tmp_handle.write(FILE_TEXT.encode('utf8')) + tmp_handle.close() + + try: + check_call('{python} {this_file} read {tmp_file}' + .format(python=sys.executable, this_file=__file__, + tmp_file=tmp_file), + shell=True) + except CalledProcessError as cpe: + self.fail(cpe.output) + + def test_uopen_redirect(self): + """ + Test redirection of unicode output to files does not raise error + + TODO: test this on non-linux OSs + """ + with temp_file(False) as (tmp_handle, tmp_file): + tmp_handle.write(FILE_TEXT.encode('utf8')) + tmp_handle.close() + + with temp_file() as redirect_file: + try: + check_call( + '{python} {this_file} read {tmp_file} >{redirect_file}' + .format(python=sys.executable, this_file=__file__, + tmp_file=tmp_file, redirect_file=redirect_file), + shell=True) + except CalledProcessError as cpe: + self.fail(cpe.output) + + @unittest.skipIf(not sys.platform.startswith('linux'), + 'Only tested on linux sofar') + def test_uopen_no_lang(self): + """ + Test that uopen in a C-LANG environment is ok + + TODO: Adapt this for other OSs; for win create batch script + """ + with temp_file(False) as (tmp_handle, tmp_file): + tmp_handle.write(FILE_TEXT.encode('utf8')) + tmp_handle.close() + + try: + check_call('LANG=C {python} {this_file} read {tmp_file}' + .format(python=sys.executable, this_file=__file__, + tmp_file=tmp_file), + shell=True) + except CalledProcessError as cpe: + self.fail(cpe.output) + + +def run_read(filename): + """This is called from test_uopen* tests as script. Reads text, compares""" + from oletools.common import uopen + # open file + with uopen(filename, 'rt') as reader: + # a few tests + if reader.closed: + raise ValueError('handle is closed!') + if reader.name != filename: + raise ValueError('Wrong filename {}'.format(reader.name)) + if reader.isatty(): + raise ValueError('Reader is a tty!') + if reader.tell() != 0: + raise ValueError('Reader.tell is not 0 at beginning') + + # read text + text = reader.read() + + # a few more tests + if not reader.closed: + raise ValueError('Reader is not closed outside context') + if reader.name != filename: + raise ValueError('Wrong filename {} after context'.format(reader.name)) + if reader.isatty(): + raise ValueError('Reader has become a tty!') + + # compare text + if sys.version_info.major <= 2: # in python2 get encoded byte string + expect = FILE_TEXT.encode('utf8') + else: # python3: should get real unicode + expect = FILE_TEXT + if text != expect: + raise ValueError('Wrong contents: {!r} != {!r}' + .format(text, expect)) + return 0 + + def run_print(): """This is called from test_read* tests as script. Prints & logs unicode""" - # hack required to import common from parent dir, not system-wide one - # (usually unittest seems to do that for us) - from os.path import abspath, dirname, join - ole_base = dirname(dirname(dirname(abspath(__file__)))) - sys.path.insert(0, ole_base) + from oletools.common import ensure_stdout_handles_unicode + from oletools.common.log_helper import log_helper + ensure_stdout_handles_unicode() + print(u'Check: \u2713') # print check mark - from oletools import common - common.ensure_stdout_handles_unicode() - print(u'\u2713') # print check mark + # check logging as well + logger = log_helper.get_or_create_silent_logger('test_encoding_handler') + log_helper.enable_logging(False, 'debug', stream=sys.stdout) + logger.info(u'Check: \u2713') + return 0 # tests call this file as script if __name__ == '__main__': if len(sys.argv) < 2: sys.exit(unittest.main()) + + # hack required to import common from parent dir, not system-wide one + # (usually unittest seems to do that for us) + from os.path import abspath, dirname, join + ole_base = dirname(dirname(dirname(abspath(__file__)))) + sys.path.insert(0, ole_base) + if sys.argv[1] == 'print': if len(sys.argv) > 2: print('Expect no arg for "print"', file=sys.stderr) sys.exit(2) sys.exit(run_print()) + elif sys.argv[1] == 'read': + if len(sys.argv) != 3: + print('Expect single arg for "read"', file=sys.stderr) + sys.exit(2) + sys.exit(run_read(sys.argv[2])) else: print('Unexpected argument: {}'.format(sys.argv[1]), file=sys.stderr) sys.exit(2)