91ca551f97b4576c680711e826a1855fb944c872
[SubU] /
1 # Copyright Jonathan Hartley 2013. BSD 3-Clause license, see LICENSE file.
2 from io import StringIO, TextIOWrapper
3 from unittest import TestCase, main
4 try:
5     from contextlib import ExitStack
6 except ImportError:
7     # python 2
8     from contextlib2 import ExitStack
9
10 try:
11     from unittest.mock import MagicMock, Mock, patch
12 except ImportError:
13     from mock import MagicMock, Mock, patch
14
15 from ..ansitowin32 import AnsiToWin32, StreamWrapper
16 from ..win32 import ENABLE_VIRTUAL_TERMINAL_PROCESSING
17 from .utils import osname
18
19
20 class StreamWrapperTest(TestCase):
21
22     def testIsAProxy(self):
23         mockStream = Mock()
24         wrapper = StreamWrapper(mockStream, None)
25         self.assertTrue( wrapper.random_attr is mockStream.random_attr )
26
27     def testDelegatesWrite(self):
28         mockStream = Mock()
29         mockConverter = Mock()
30         wrapper = StreamWrapper(mockStream, mockConverter)
31         wrapper.write('hello')
32         self.assertTrue(mockConverter.write.call_args, (('hello',), {}))
33
34     def testDelegatesContext(self):
35         mockConverter = Mock()
36         s = StringIO()
37         with StreamWrapper(s, mockConverter) as fp:
38             fp.write(u'hello')
39         self.assertTrue(s.closed)
40
41     def testProxyNoContextManager(self):
42         mockStream = MagicMock()
43         mockStream.__enter__.side_effect = AttributeError()
44         mockConverter = Mock()
45         with self.assertRaises(AttributeError) as excinfo:
46             with StreamWrapper(mockStream, mockConverter) as wrapper:
47                 wrapper.write('hello')
48
49     def test_closed_shouldnt_raise_on_closed_stream(self):
50         stream = StringIO()
51         stream.close()
52         wrapper = StreamWrapper(stream, None)
53         self.assertEqual(wrapper.closed, True)
54
55     def test_closed_shouldnt_raise_on_detached_stream(self):
56         stream = TextIOWrapper(StringIO())
57         stream.detach()
58         wrapper = StreamWrapper(stream, None)
59         self.assertEqual(wrapper.closed, True)
60
61 class AnsiToWin32Test(TestCase):
62
63     def testInit(self):
64         mockStdout = Mock()
65         auto = Mock()
66         stream = AnsiToWin32(mockStdout, autoreset=auto)
67         self.assertEqual(stream.wrapped, mockStdout)
68         self.assertEqual(stream.autoreset, auto)
69
70     @patch('colorama.ansitowin32.winterm', None)
71     @patch('colorama.ansitowin32.winapi_test', lambda *_: True)
72     def testStripIsTrueOnWindows(self):
73         with osname('nt'):
74             mockStdout = Mock()
75             stream = AnsiToWin32(mockStdout)
76             self.assertTrue(stream.strip)
77
78     def testStripIsFalseOffWindows(self):
79         with osname('posix'):
80             mockStdout = Mock(closed=False)
81             stream = AnsiToWin32(mockStdout)
82             self.assertFalse(stream.strip)
83
84     def testWriteStripsAnsi(self):
85         mockStdout = Mock()
86         stream = AnsiToWin32(mockStdout)
87         stream.wrapped = Mock()
88         stream.write_and_convert = Mock()
89         stream.strip = True
90
91         stream.write('abc')
92
93         self.assertFalse(stream.wrapped.write.called)
94         self.assertEqual(stream.write_and_convert.call_args, (('abc',), {}))
95
96     def testWriteDoesNotStripAnsi(self):
97         mockStdout = Mock()
98         stream = AnsiToWin32(mockStdout)
99         stream.wrapped = Mock()
100         stream.write_and_convert = Mock()
101         stream.strip = False
102         stream.convert = False
103
104         stream.write('abc')
105
106         self.assertFalse(stream.write_and_convert.called)
107         self.assertEqual(stream.wrapped.write.call_args, (('abc',), {}))
108
109     def assert_autoresets(self, convert, autoreset=True):
110         stream = AnsiToWin32(Mock())
111         stream.convert = convert
112         stream.reset_all = Mock()
113         stream.autoreset = autoreset
114         stream.winterm = Mock()
115
116         stream.write('abc')
117
118         self.assertEqual(stream.reset_all.called, autoreset)
119
120     def testWriteAutoresets(self):
121         self.assert_autoresets(convert=True)
122         self.assert_autoresets(convert=False)
123         self.assert_autoresets(convert=True, autoreset=False)
124         self.assert_autoresets(convert=False, autoreset=False)
125
126     def testWriteAndConvertWritesPlainText(self):
127         stream = AnsiToWin32(Mock())
128         stream.write_and_convert( 'abc' )
129         self.assertEqual( stream.wrapped.write.call_args, (('abc',), {}) )
130
131     def testWriteAndConvertStripsAllValidAnsi(self):
132         stream = AnsiToWin32(Mock())
133         stream.call_win32 = Mock()
134         data = [
135             'abc\033[mdef',
136             'abc\033[0mdef',
137             'abc\033[2mdef',
138             'abc\033[02mdef',
139             'abc\033[002mdef',
140             'abc\033[40mdef',
141             'abc\033[040mdef',
142             'abc\033[0;1mdef',
143             'abc\033[40;50mdef',
144             'abc\033[50;30;40mdef',
145             'abc\033[Adef',
146             'abc\033[0Gdef',
147             'abc\033[1;20;128Hdef',
148         ]
149         for datum in data:
150             stream.wrapped.write.reset_mock()
151             stream.write_and_convert( datum )
152             self.assertEqual(
153                [args[0] for args in stream.wrapped.write.call_args_list],
154                [ ('abc',), ('def',) ]
155             )
156
157     def testWriteAndConvertSkipsEmptySnippets(self):
158         stream = AnsiToWin32(Mock())
159         stream.call_win32 = Mock()
160         stream.write_and_convert( '\033[40m\033[41m' )
161         self.assertFalse( stream.wrapped.write.called )
162
163     def testWriteAndConvertCallsWin32WithParamsAndCommand(self):
164         stream = AnsiToWin32(Mock())
165         stream.convert = True
166         stream.call_win32 = Mock()
167         stream.extract_params = Mock(return_value='params')
168         data = {
169             'abc\033[adef':         ('a', 'params'),
170             'abc\033[;;bdef':       ('b', 'params'),
171             'abc\033[0cdef':        ('c', 'params'),
172             'abc\033[;;0;;Gdef':    ('G', 'params'),
173             'abc\033[1;20;128Hdef': ('H', 'params'),
174         }
175         for datum, expected in data.items():
176             stream.call_win32.reset_mock()
177             stream.write_and_convert( datum )
178             self.assertEqual( stream.call_win32.call_args[0], expected )
179
180     def test_reset_all_shouldnt_raise_on_closed_orig_stdout(self):
181         stream = StringIO()
182         converter = AnsiToWin32(stream)
183         stream.close()
184
185         converter.reset_all()
186
187     def test_wrap_shouldnt_raise_on_closed_orig_stdout(self):
188         stream = StringIO()
189         stream.close()
190         with \
191             patch("colorama.ansitowin32.os.name", "nt"), \
192             patch("colorama.ansitowin32.winapi_test", lambda: True):
193                 converter = AnsiToWin32(stream)
194         self.assertTrue(converter.strip)
195         self.assertFalse(converter.convert)
196
197     def test_wrap_shouldnt_raise_on_missing_closed_attr(self):
198         with \
199             patch("colorama.ansitowin32.os.name", "nt"), \
200             patch("colorama.ansitowin32.winapi_test", lambda: True):
201                 converter = AnsiToWin32(object())
202         self.assertTrue(converter.strip)
203         self.assertFalse(converter.convert)
204
205     def testExtractParams(self):
206         stream = AnsiToWin32(Mock())
207         data = {
208             '':               (0,),
209             ';;':             (0,),
210             '2':              (2,),
211             ';;002;;':        (2,),
212             '0;1':            (0, 1),
213             ';;003;;456;;':   (3, 456),
214             '11;22;33;44;55': (11, 22, 33, 44, 55),
215         }
216         for datum, expected in data.items():
217             self.assertEqual(stream.extract_params('m', datum), expected)
218
219     def testCallWin32UsesLookup(self):
220         listener = Mock()
221         stream = AnsiToWin32(listener)
222         stream.win32_calls = {
223             1: (lambda *_, **__: listener(11),),
224             2: (lambda *_, **__: listener(22),),
225             3: (lambda *_, **__: listener(33),),
226         }
227         stream.call_win32('m', (3, 1, 99, 2))
228         self.assertEqual(
229             [a[0][0] for a in listener.call_args_list],
230             [33, 11, 22] )
231
232     def test_osc_codes(self):
233         mockStdout = Mock()
234         stream = AnsiToWin32(mockStdout, convert=True)
235         with patch('colorama.ansitowin32.winterm') as winterm:
236             data = [
237                 '\033]0\x07',                      # missing arguments
238                 '\033]0;foo\x08',                  # wrong OSC command
239                 '\033]0;colorama_test_title\x07',  # should work
240                 '\033]1;colorama_test_title\x07',  # wrong set command
241                 '\033]2;colorama_test_title\x07',  # should work
242                 '\033]' + ';' * 64 + '\x08',       # see issue #247
243             ]
244             for code in data:
245                 stream.write(code)
246             self.assertEqual(winterm.set_title.call_count, 2)
247
248     def test_native_windows_ansi(self):
249         with ExitStack() as stack:
250             def p(a, b):
251                 stack.enter_context(patch(a, b, create=True))
252             # Pretend to be on Windows
253             p("colorama.ansitowin32.os.name", "nt")
254             p("colorama.ansitowin32.winapi_test", lambda: True)
255             p("colorama.win32.winapi_test", lambda: True)
256             p("colorama.winterm.win32.windll", "non-None")
257             p("colorama.winterm.get_osfhandle", lambda _: 1234)
258
259             # Pretend that our mock stream has native ANSI support
260             p(
261                 "colorama.winterm.win32.GetConsoleMode",
262                 lambda _: ENABLE_VIRTUAL_TERMINAL_PROCESSING,
263             )
264             SetConsoleMode = Mock()
265             p("colorama.winterm.win32.SetConsoleMode", SetConsoleMode)
266
267             stdout = Mock()
268             stdout.closed = False
269             stdout.isatty.return_value = True
270             stdout.fileno.return_value = 1
271
272             # Our fake console says it has native vt support, so AnsiToWin32 should
273             # enable that support and do nothing else.
274             stream = AnsiToWin32(stdout)
275             SetConsoleMode.assert_called_with(1234, ENABLE_VIRTUAL_TERMINAL_PROCESSING)
276             self.assertFalse(stream.strip)
277             self.assertFalse(stream.convert)
278             self.assertFalse(stream.should_wrap())
279
280             # Now let's pretend we're on an old Windows console, that doesn't have
281             # native ANSI support.
282             p("colorama.winterm.win32.GetConsoleMode", lambda _: 0)
283             SetConsoleMode = Mock()
284             p("colorama.winterm.win32.SetConsoleMode", SetConsoleMode)
285
286             stream = AnsiToWin32(stdout)
287             SetConsoleMode.assert_called_with(1234, ENABLE_VIRTUAL_TERMINAL_PROCESSING)
288             self.assertTrue(stream.strip)
289             self.assertTrue(stream.convert)
290             self.assertTrue(stream.should_wrap())
291
292
293 if __name__ == '__main__':
294     main()