Skip to content

Improve tests for _colorize.can_colorize() #129234

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 5 commits into from
Jan 24, 2025
Merged

Improve tests for _colorize.can_colorize() #129234

Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
8e64217
Improve tests for _colorize.can_colorize()
serhiy-storchaka Jan 23, 2025
3b870a6
Fix test on Windows.
serhiy-storchaka Jan 23, 2025
d7bac3f
Merge branch 'main' into test__colorize
serhiy-storchaka Jan 23, 2025
be47b1b
Fix a test again.
serhiy-storchaka Jan 23, 2025
5c2c829
Update Lib/test/test__colorize.py
serhiy-storchaka Jan 24, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
164 changes: 108 additions & 56 deletions Lib/test/test__colorize.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,80 +1,132 @@
import contextlib
import io
import sys
import unittest
import unittest.mock
import _colorize
from test.support import force_not_colorized, make_clean_env
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

make_clean_env was moved to test.support from the PyREPL tests for use in this file (in #127877).

If we're not using it here any more, should we move it back, or leave it in case it's useful for others later?

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I do not know. I tried to use make_clean_env(), but we need to remove also NO_COLOR. On other hand, the only PYTHON* variable that needs to be removed is PYTHON_COLORS. Let it be there for now, and then we'll see.

from test.support.os_helper import EnvironmentVarGuard

ORIGINAL_CAN_COLORIZE = _colorize.can_colorize

@contextlib.contextmanager
def clear_env():
with EnvironmentVarGuard() as mock_env:
for var in "FORCE_COLOR", "NO_COLOR", "PYTHON_COLORS":
mock_env.unset(var)
yield mock_env

def setUpModule():
_colorize.can_colorize = lambda *args, **kwargs: False


def tearDownModule():
_colorize.can_colorize = ORIGINAL_CAN_COLORIZE
def supports_virtual_terminal():
if sys.platform == "win32":
return unittest.mock.patch("nt._supports_virtual_terminal", return_value=True)
else:
return contextlib.nullcontext()


class TestColorizeFunction(unittest.TestCase):
def setUp(self):
# Remove PYTHON* environment variables to isolate from local user
# settings and simulate running with `-E`. Such variables should be
# added to test methods later to patched os.environ.
patcher = unittest.mock.patch("os.environ", new=make_clean_env())
self.addCleanup(patcher.stop)
patcher.start()

@force_not_colorized
def test_colorized_detection_checks_for_environment_variables(self):
flags = unittest.mock.MagicMock(ignore_environment=False)
def check(env, fallback, expected):
with (self.subTest(env=env, fallback=fallback),
clear_env() as mock_env):
mock_env.update(env)
isatty_mock.return_value = fallback
stdout_mock.isatty.return_value = fallback
self.assertEqual(_colorize.can_colorize(), expected)

with (unittest.mock.patch("os.isatty") as isatty_mock,
unittest.mock.patch("sys.stdout") as stdout_mock,
unittest.mock.patch("sys.stderr") as stderr_mock,
unittest.mock.patch("sys.flags", flags),
unittest.mock.patch("_colorize.can_colorize", ORIGINAL_CAN_COLORIZE),
(unittest.mock.patch("nt._supports_virtual_terminal", return_value=False)
if sys.platform == "win32" else
contextlib.nullcontext()) as vt_mock):
supports_virtual_terminal()):
stdout_mock.fileno.return_value = 1

for fallback in False, True:
check({}, fallback, fallback)
check({'TERM': 'dumb'}, fallback, False)
check({'TERM': 'xterm'}, fallback, fallback)
check({'TERM': ''}, fallback, fallback)
check({'FORCE_COLOR': '1'}, fallback, True)
check({'FORCE_COLOR': '0'}, fallback, True)
check({'NO_COLOR': '1'}, fallback, False)
check({'NO_COLOR': '0'}, fallback, False)

check({'TERM': 'dumb', 'FORCE_COLOR': '1'}, False, True)
check({'FORCE_COLOR': '1', 'NO_COLOR': '1'}, True, False)

for ignore_environment in False, True:
# Simulate running with or without `-E`.
flags = unittest.mock.MagicMock(ignore_environment=ignore_environment)
with unittest.mock.patch("sys.flags", flags):
check({'PYTHON_COLORS': '1'}, True, True)
check({'PYTHON_COLORS': '1'}, False, not ignore_environment)
check({'PYTHON_COLORS': '0'}, True, ignore_environment)
check({'PYTHON_COLORS': '0'}, False, False)
for fallback in False, True:
check({'PYTHON_COLORS': 'x'}, fallback, fallback)
check({'PYTHON_COLORS': ''}, fallback, fallback)

check({'TERM': 'dumb', 'PYTHON_COLORS': '1'}, False, not ignore_environment)
check({'NO_COLOR': '1', 'PYTHON_COLORS': '1'}, False, not ignore_environment)
check({'FORCE_COLOR': '1', 'PYTHON_COLORS': '0'}, True, ignore_environment)

@unittest.skipUnless(sys.platform == "win32", "requires Windows")
def test_colorized_detection_checks_on_windows(self):
with (clear_env(),
unittest.mock.patch("os.isatty") as isatty_mock,
unittest.mock.patch("sys.stdout") as stdout_mock,
supports_virtual_terminal() as vt_mock):
stdout_mock.fileno.return_value = 1
isatty_mock.return_value = True
stdout_mock.isatty.return_value = True

vt_mock.return_value = True
self.assertEqual(_colorize.can_colorize(), True)
vt_mock.return_value = False
self.assertEqual(_colorize.can_colorize(), False)
import nt
del nt._supports_virtual_terminal
self.assertEqual(_colorize.can_colorize(), False)

def test_colorized_detection_checks_for_std_streams(self):
with (clear_env(),
unittest.mock.patch("os.isatty") as isatty_mock,
unittest.mock.patch("sys.stdout") as stdout_mock,
unittest.mock.patch("sys.stderr") as stderr_mock,
supports_virtual_terminal()):
stdout_mock.fileno.return_value = 1
stderr_mock.fileno.side_effect = ZeroDivisionError
stderr_mock.isatty.side_effect = ZeroDivisionError

isatty_mock.return_value = True
stdout_mock.isatty.return_value = True
stderr_mock.fileno.return_value = 2
stderr_mock.isatty.return_value = True
with unittest.mock.patch("os.environ", {'TERM': 'dumb'}):
self.assertEqual(_colorize.can_colorize(), False)
with unittest.mock.patch("os.environ", {'PYTHON_COLORS': '1'}):
self.assertEqual(_colorize.can_colorize(), True)
with unittest.mock.patch("os.environ", {'PYTHON_COLORS': '0'}):
self.assertEqual(_colorize.can_colorize(), False)
with unittest.mock.patch("os.environ", {'NO_COLOR': '1'}):
self.assertEqual(_colorize.can_colorize(), False)
with unittest.mock.patch("os.environ",
{'NO_COLOR': '1', "PYTHON_COLORS": '1'}):
self.assertEqual(_colorize.can_colorize(), True)
with unittest.mock.patch("os.environ", {'FORCE_COLOR': '1'}):
self.assertEqual(_colorize.can_colorize(), True)
with unittest.mock.patch("os.environ",
{'FORCE_COLOR': '1', 'NO_COLOR': '1'}):
self.assertEqual(_colorize.can_colorize(), False)
with unittest.mock.patch("os.environ",
{'FORCE_COLOR': '1', "PYTHON_COLORS": '0'}):
self.assertEqual(_colorize.can_colorize(), False)

with unittest.mock.patch("os.environ", {}):
if sys.platform == "win32":
self.assertEqual(_colorize.can_colorize(), False)

vt_mock.return_value = True
self.assertEqual(_colorize.can_colorize(), True)
else:
self.assertEqual(_colorize.can_colorize(), True)
self.assertEqual(_colorize.can_colorize(), True)

isatty_mock.return_value = False
stdout_mock.isatty.return_value = False
self.assertEqual(_colorize.can_colorize(), False)

def test_colorized_detection_checks_for_file(self):
with clear_env(), supports_virtual_terminal():

with unittest.mock.patch("os.isatty") as isatty_mock:
file = unittest.mock.MagicMock()
file.fileno.return_value = 1
isatty_mock.return_value = True
self.assertEqual(_colorize.can_colorize(file=file), True)
isatty_mock.return_value = False
stdout_mock.isatty.return_value = False
stderr_mock.isatty.return_value = False
self.assertEqual(_colorize.can_colorize(), False)
self.assertEqual(_colorize.can_colorize(file=file), False)

# No file.fileno.
with unittest.mock.patch("os.isatty", side_effect=ZeroDivisionError):
file = unittest.mock.MagicMock(spec=['isatty'])
file.isatty.return_value = True
self.assertEqual(_colorize.can_colorize(file=file), False)

# file.fileno() raises io.UnsupportedOperation.
with unittest.mock.patch("os.isatty", side_effect=ZeroDivisionError):
file = unittest.mock.MagicMock()
file.fileno.side_effect = io.UnsupportedOperation
file.isatty.return_value = True
self.assertEqual(_colorize.can_colorize(file=file), True)
file.isatty.return_value = False
self.assertEqual(_colorize.can_colorize(file=file), False)


if __name__ == "__main__":
Expand Down
Loading