test_encoding_handler.py
7.6 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
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
"""Test common.ensure_stdout_handles_unicode"""
from __future__ import print_function
import unittest
import sys
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.
If `just_name` is `False` this yields (file-name, open-file-handle),
if `just_name` is `True` this yields just the file-name (and closes
the file-handle if we are on windows)
"""
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 ... except for Windows where it is not possible for another
# process to write to that handle
if just_name:
if sys.platform.startswith('win'):
os.close(tmp_descriptor)
tmp_descriptor = None
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_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_print_redirect(self):
"""Test redirection of unicode output to files does not raise error."""
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)
@unittest.skipIf(not sys.platform.startswith('linux'),
'Need to adapt this test to Windows')
def test_print_no_lang(self):
"""
Test redirection of unicode output to files does not raise error
TODO: Adapt this for other OSs; for win create batch script
"""
check_call('LANG=C {python} {this_file} print'
.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."""
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'),
'Need to adapt this test to Windows')
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.io_encoding 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))
# the following test raises an exception because reader is closed, so isatty cannot be called:
# 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"""
from oletools.common.io_encoding import ensure_stdout_handles_unicode
from oletools.common.log_helper import log_helper
ensure_stdout_handles_unicode()
print(u'Check: \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)