8000 Improve tests for _colorize.can_colorize() (#129234) · python/cpython@7907203 · GitHub
[go: up one dir, main page]
Skip to content

Commit 7907203

Browse files
serhiy-storchakahugovk
serhiy-storchaka
and
hugovk
authored
Improve tests for _colorize.can_colorize() (#129234)
Co-authored-by: Hugo van Kemenade <1324225+hugovk@users.noreply.github.com>
1 parent 8e0b360 commit 7907203

File tree

1 file changed

+108
-56
lines changed

1 file changed

+108
-56
lines changed

‎Lib/test/test__colorize.py

Lines changed: 108 additions & 56 deletions
Original file line numberDiff line numberDiff line change
@@ -1,80 +1,132 @@
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
1317

14-
15-
def tearDownModule():
16-
_colorize.can_colorize = ORIGINAL_CAN_COLORIZE
18+
def supports_virtual_terminal():
19+
if sys.platform == "win32":
20+
return unittest.mock.patch("nt._supports_virtual_terminal", return_value=True)
21+
else:
22+
return contextlib.nullcontext()
1723

1824

1925
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
2926
def test_colorized_detection_checks_for_environment_variables(self):
30-
flags = unittest.mock.MagicMock(ignore_environment=False)
27+
def check(env, fallback, expected):
28+
with (self.subTest(env=env, fallback=fallback),
29+
clear_env() as mock_env):
30+
mock_env.update(env)
31+
isatty_mock.return_value = fallback
32+
stdout_mock.isatty.return_value = fallback
33+
self.assertEqual(_colorize.can_colorize(), expected)
34+
3135
with (unittest.mock.patch("os.isatty") as isatty_mock,
3236
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):
37+
supports_virtual_terminal()):
38+
stdout_mock.fileno.return_value = 1
39+
40+
for fallback in False, True:
41+
check({}, fallback, fallback)
42+
check({'TERM': 'dumb'}, fallback, False)
43+
check({'TERM': 'xterm'}, fallback, fallback)
44+
check({'TERM': ''}, fallback, fallback)
45+
check({'FORCE_COLOR': '1'}, fallback, True)
46+
check({'FORCE_COLOR': '0'}, fallback, True)
47+
check({'NO_COLOR': '1'}, fallback, False)
48+
check({'NO_COLOR': '0'}, fallback, False)
49+
50+
check({'TERM': 'dumb', 'FORCE_COLOR': '1'}, False, True)
51+
check({'FORCE_COLOR': '1', 'NO_COLOR': '1'}, True, False)
3952

53+
for ignore_environment in False, True:
54+
# Simulate running with or without `-E`.
55+
flags = unittest.mock.MagicMock(ignore_environment=ignore_environment)
56+
with unittest.mock.patch("sys.flags", flags):
57+
check({'PYTHON_COLORS': '1'}, True, True)
58+
check({'PYTHON_COLORS': '1'}, False, not ignore_environment)
59+
check({'PYTHON_COLORS': '0'}, True, ignore_environment)
60+
check({'PYTHON_COLORS': '0'}, False, False)
61+
for fallback in False, True:
62+
check({'PYTHON_COLORS': 'x'}, fallback, fallback)
63+
check({'PYTHON_COLORS': ''}, fallback, fallback)
64+
65+
check({'TERM': 'dumb', 'PYTHON_COLORS': '1'}, False, not ignore_environment)
66+
check({'NO_COLOR': '1', 'PYTHON_COLORS': '1'}, False, not ignore_environment)
67+
check({'FORCE_COLOR': '1', 'PYTHON_COLORS': '0'}, True, ignore_environment)
68+
69+
@unittest.skipUnless(sys.platform == "win32", "requires Windows")
70+
def test_colorized_detection_checks_on_windows(self):
71+
with (clear_env(),
72+
unittest.mock.patch("os.isatty") as isatty_mock,
73+
unittest.mock.patch("sys.stdout") as stdout_mock,
74+
supports_virtual_terminal() as vt_mock):
75+
stdout_mock.fileno.return_value = 1
4076
isatty_mock.return_value = True
77+
stdout_mock.isatty.return_value = True
78+
79+
vt_mock.return_value = True
80+
self.assertEqual(_colorize.can_colorize(), True)
81+
vt_mock.return_value = False
82+
self.assertEqual(_colorize.can_colorize(), False)
83+
import nt
84+
del nt._supports_virtual_terminal
85+
self.assertEqual(_colorize.can_colorize(), False)
86+
87+
def test_colorized_detection_checks_for_std_streams(self):
88+
with (clear_env(),
89+
unittest.mock.patch("os.isatty") as isatty_mock,
90+
unittest.mock.patch("sys.stdout") as stdout_mock,
91+
unittest.mock.patch("sys.stderr") as stderr_mock,
92+
supports_virtual_terminal()):
4193
stdout_mock.fileno.return_value = 1
94+
stderr_mock.fileno.side_effect = ZeroDivisionError
95+
stderr_mock.isatty.side_effect = ZeroDivisionError
96+
97+
isatty_mock.return_value = True
4298
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)
99+
self.assertEqual(_colorize.can_colorize(), True)
100+
101+
isatty_mock.return_value = False
102+
stdout_mock.isatty.return_value = False
103+
self.assertEqual(_colorize.can_colorize(), False)
104+
105+
def test_colorized_detection_checks_for_file(self):
106+
with clear_env(), supports_virtual_terminal():
73107

108+
with unittest.mock.patch("os.isatty") as isatty_mock:
109+
file = unittest.mock.MagicMock()
110+
file.fileno.return_value = 1
111+
isatty_mock.return_value = True
112+
self.assertEqual(_colorize.can_colorize(file=file), True)
74113
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)
114+
self.assertEqual(_colorize.can_colorize(file=file), False)
115+
116+
# No file.fileno.
117+
with unittest.mock.patch("os.isatty", side_effect=ZeroDivisionError):
118+
file = unittest.mock.MagicMock(spec=['isatty'])
119+
file.isatty.return_value = True
120+
self.assertEqual(_colorize.can_colorize(file=file), False)
121+
122+
# file.fileno() raises io.UnsupportedOperation.
123+
with unittest.mock.patch("os.isatty", side_effect=ZeroDivisionError):
124+
file = unittest.mock.MagicMock()
125+
file.fileno.side_effect = io.UnsupportedOperation
126+
file.isatty.return_value = True
127+
self.assertEqual(_colorize.can_colorize(file=file), True)
128+
file.isatty.return_value = False
129+
self.assertEqual(_colorize.can_colorize(file=file), False)
78130

79131

80132
if __name__ == "__main__":

0 commit comments

Comments
 (0)
0