Improve tests for _colorize.can_colorize() (#129234)

Co-authored-by: Hugo van Kemenade <1324225+hugovk@users.noreply.github.com>

diff --git a/Lib/test/test__colorize.py b/Lib/test/test__colorize.py
index 25519ba7e92e1fca5c059e1f5de927fc381c4ba1..ff15d83cb078fef474379c16c64d8aaffe32c650 100644 (file)
--- a/Lib/test/test__colorize.py
+++ b/Lib/test/test__colorize.py
@@ -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
+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__":