Skip to content

Commit 8e64217

Browse files
serhiy-storchakaserhiy-storchaka
committed
Improve tests for _colorize.can_colorize()
1 parent cf0b2da commit 8e64217

File tree

1 file changed

+107
-57
lines changed

1 file changed

+107
-57
lines changed

‎Lib/test/test__colorize.py‎

Lines changed: 107 additions & 57 deletions
Original file line numberDiff line numberDiff line change
@@ -1,80 +1,130 @@
11
import contextlib
2+
import io
23
import sys
34
import unittest
45
import unittest.mock
56
import _colorize
6-
from test.support import force_not_colorized, make_clean_env
7+
from test.support.os_helper import EnvironmentVarGuard
78

8-
ORIGINAL_CAN_COLORIZE = _colorize.can_colorize
99

10+
@contextlib.contextmanager
11+
def clear_env():
12+
with EnvironmentVarGuard() as mock_env:
13+
for var in "FORCE_COLOR", "NO_COLOR", "PYTHON_COLORS":
14+
mock_env.unset(var)
15+
yield mock_env
1016

11-
def setUpModule():
12-
_colorize.can_colorize = lambda *args, **kwargs: False
13-
14-
15-
def tearDownModule():
16-
_colorize.can_colorize = ORIGINAL_CAN_COLORIZE
17+
def supports_virtual_terminal():
18+
if sys.platform == "win32":
19+
return unittest.mock.patch("nt._supports_virtual_terminal", return_value=True)
20+
else:
21+
return contextlib.nullcontext()
1722

1823

1924
class TestColorizeFunction(unittest.TestCase):
20-
def setUp(self):
21-
# Remove PYTHON* environment variables to isolate from local user
22-
# settings and simulate running with `-E`. Such variables should be
23-
# added to test methods later to patched os.environ.
24-
patcher = unittest.mock.patch("os.environ", new=make_clean_env())
25-
self.addCleanup(patcher.stop)
26-
patcher.start()
27-
28-
@force_not_colorized
2925
def test_colorized_detection_checks_for_environment_variables(self):
30-
flags = unittest.mock.MagicMock(ignore_environment=False)
26+
def check(env, fallback, expected):
27+
with (self.subTest(env=env, fallback=fallback),
28+
clear_env() as mock_env):
29+
mock_env.update(env)
30+
isatty_mock.return_value = fallback
31+
stdout_mock.isatty.return_value = fallback
32+
self.assertEqual(_colorize.can_colorize(), expected)
33+
3134
with (unittest.mock.patch("os.isatty") as isatty_mock,
3235
unittest.mock.patch("sys.stdout") as stdout_mock,
33-
unittest.mock.patch("sys.stderr") as stderr_mock,
34-
unittest.mock.patch("sys.flags", flags),
35-
unittest.mock.patch("_colorize.can_colorize", ORIGINAL_CAN_COLORIZE),
36-
(unittest.mock.patch("nt._supports_virtual_terminal", return_value=False)
37-
if sys.platform == "win32" else
38-
contextlib.nullcontext()) as vt_mock):
36+
supports_virtual_terminal()):
37+
stdout_mock.fileno.return_value = 1
38+
39+
for fallback in False, True:
40+
check({}, fallback, fallback)
41+
check({'TERM': 'dumb'}, fallback, False)
42+
check({'TERM': 'xterm'}, fallback, fallback)
43+
check({'TERM': ''}, fallback, fallback)
44+
check({'FORCE_COLOR': '1'}, fallback, True)
45+
check({'FORCE_COLOR': '0'}, fallback, True)
46+
check({'NO_COLOR': '1'}, fallback, False)
47+
check({'NO_COLOR': '0'}, fallback, False)
3948

49+
check({'TERM': 'dumb', 'FORCE_COLOR': '1'}, False, True)
50+
check({'FORCE_COLOR': '1', 'NO_COLOR': '1'}, True, False)
51+
52+
for ignore_environment in False, True:
53+
# Simulate running with or without `-E`.
54+
flags = unittest.mock.MagicMock(ignore_environment=ignore_environment)
55+
with unittest.mock.patch("sys.flags", flags):
56+
check({'PYTHON_COLORS': '1'}, True, True)
57+
check({'PYTHON_COLORS': '1'}, False, not ignore_environment)
58+
check({'PYTHON_COLORS': '0'}, True, ignore_environment)
59+
check({'PYTHON_COLORS': '0'}, False, False)
60+
for fallback in False, True:
61+
check({'PYTHON_COLORS': 'x'}, fallback, fallback)
62+
check({'PYTHON_COLORS': ''}, fallback, fallback)
63+
64+
check({'TERM': 'dumb', 'PYTHON_COLORS': '1'}, False, not ignore_environment)
65+
check({'NO_COLOR': '1', 'PYTHON_COLORS': '1'}, False, not ignore_environment)
66+
check({'FORCE_COLOR': '1', 'PYTHON_COLORS': '0'}, True, ignore_environment)
67+
68+
@unittest.skipUnless(sys.platform == "win32", "requires Windows")
69+
def test_colorized_detection_checks_on_windows(self):
70+
with (clear_env(),
71+
unittest.mock.patch("os.isatty") as isatty_mock,
72+
unittest.mock.patch("sys.stdout") as stdout_mock,
73+
supports_virtual_terminal() as vt_mock):
74+
stdout_mock.fileno.return_value = 1
4075
isatty_mock.return_value = True
76+
stdout_mock.isatty.return_value = True
77+
78+
vt_mock.return_value = True
79+
self.assertEqual(_colorize.can_colorize(), True)
80+
vt_mock.return_value = False
81+
self.assertEqual(_colorize.can_colorize(), True)
82+
del nt._supports_virtual_terminal
83+
self.assertEqual(_colorize.can_colorize(), True)
84+
85+
def test_colorized_detection_checks_for_std_streams(self):
86+
with (clear_env(),
87+
unittest.mock.patch("os.isatty") as isatty_mock,
88+
unittest.mock.patch("sys.stdout") as stdout_mock,
89+
unittest.mock.patch("sys.stderr") as stderr_mock,
90+
supports_virtual_terminal()):
4191
stdout_mock.fileno.return_value = 1
92+
stderr_mock.fileno.side_effect = ZeroDivisionError
93+
stderr_mock.isatty.side_effect = ZeroDivisionError
94+
95+
isatty_mock.return_value = True
4296
stdout_mock.isatty.return_value = True
43-
stderr_mock.fileno.return_value = 2
44-
stderr_mock.isatty.return_value = True
45-
with unittest.mock.patch("os.environ", {'TERM': 'dumb'}):
46-
self.assertEqual(_colorize.can_colorize(), False)
47-
with unittest.mock.patch("os.environ", {'PYTHON_COLORS': '1'}):
48-
self.assertEqual(_colorize.can_colorize(), True)
49-
with unittest.mock.patch("os.environ", {'PYTHON_COLORS': '0'}):
50-
self.assertEqual(_colorize.can_colorize(), False)
51-
with unittest.mock.patch("os.environ", {'NO_COLOR': '1'}):
52-
self.assertEqual(_colorize.can_colorize(), False)
53-
with unittest.mock.patch("os.environ",
54-
{'NO_COLOR': '1', "PYTHON_COLORS": '1'}):
55-
self.assertEqual(_colorize.can_colorize(), True)
56-
with unittest.mock.patch("os.environ", {'FORCE_COLOR': '1'}):
57-
self.assertEqual(_colorize.can_colorize(), True)
58-
with unittest.mock.patch("os.environ",
59-
{'FORCE_COLOR': '1', 'NO_COLOR': '1'}):
60-
self.assertEqual(_colorize.can_colorize(), False)
61-
with unittest.mock.patch("os.environ",
62-
{'FORCE_COLOR': '1', "PYTHON_COLORS": '0'}):
63-
self.assertEqual(_colorize.can_colorize(), False)
64-
65-
with unittest.mock.patch("os.environ", {}):
66-
if sys.platform == "win32":
67-
self.assertEqual(_colorize.can_colorize(), False)
68-
69-
vt_mock.return_value = True
70-
self.assertEqual(_colorize.can_colorize(), True)
71-
else:
72-
self.assertEqual(_colorize.can_colorize(), True)
97+
self.assertEqual(_colorize.can_colorize(), True)
98+
99+
isatty_mock.return_value = False
100+
stdout_mock.isatty.return_value = False
101+
self.assertEqual(_colorize.can_colorize(), False)
102+
103+
def test_colorized_detection_checks_for_file(self):
104+
with clear_env(), supports_virtual_terminal():
73105

106+
with unittest.mock.patch("os.isatty") as isatty_mock:
107+
file = unittest.mock.MagicMock()
108+
file.fileno.return_value = 1
109+
isatty_mock.return_value = True
110+
self.assertEqual(_colorize.can_colorize(file=file), True)
74111
isatty_mock.return_value = False
75-
stdout_mock.isatty.return_value = False
76-
stderr_mock.isatty.return_value = False
77-
self.assertEqual(_colorize.can_colorize(), False)
112+
self.assertEqual(_colorize.can_colorize(file=file), False)
113+
114+
# No file.fileno.
115+
with unittest.mock.patch("os.isatty", side_effect=ZeroDivisionError):
116+
file = unittest.mock.MagicMock(spec=['isatty'])
117+
file.isatty.return_value = True
118+
self.assertEqual(_colorize.can_colorize(file=file), False)
119+
120+
# file.fileno() raises io.UnsupportedOperation.
121+
with unittest.mock.patch("os.isatty", side_effect=ZeroDivisionError):
122+
file = unittest.mock.MagicMock()
123+
file.fileno.side_effect = io.UnsupportedOperation
124+
file.isatty.return_value = True
125+
self.assertEqual(_colorize.can_colorize(file=file), True)
126+
file.isatty.return_value = False
127+
self.assertEqual(_colorize.can_colorize(file=file), False)
78128

79129

80130
if __name__ == "__main__":

0 commit comments

Comments
 (0)