Refactor colorize tests

PalmtopTiger · PalmtopTiger · commit 6946167feeed · 2024-12-14T13:06:23.000+07:00
diff --git a/Lib/test/test__colorize.py b/Lib/test/test__colorize.py
@@ -19,57 +19,15 @@ def tearDownModule():
 class TestColorizeFunction(unittest.TestCase):
     @force_not_colorized
     def test_colorized_detection_checks_for_environment_variables(self):
-        if sys.platform == "win32":
-            virtual_patching = unittest.mock.patch("nt._supports_virtual_terminal",
-                                                   return_value=True)
-        else:
-            virtual_patching = contextlib.nullcontext()
-        with virtual_patching:
-
-            flags = unittest.mock.MagicMock(ignore_environment=False)
-            with (unittest.mock.patch("os.isatty") as isatty_mock,
-                  unittest.mock.patch("sys.stderr") as stderr_mock,
-                  unittest.mock.patch("sys.flags", flags),
-                  unittest.mock.patch("_colorize.can_colorize", ORIGINAL_CAN_COLORIZE)):
-                isatty_mock.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", {}):
-                    self.assertEqual(_colorize.can_colorize(), True)
-
-                isatty_mock.return_value = False
-                stderr_mock.isatty.return_value = False
-                with unittest.mock.patch("os.environ", {}):
-                    self.assertEqual(_colorize.can_colorize(), False)
-
-    @force_not_colorized
-    @unittest.skipUnless(sys.platform == "win32", "Windows only")
-    def test_colorized_detection_checks_for_environment_variables_no_vt(self):
         flags = unittest.mock.MagicMock(ignore_environment=False)
-        with (unittest.mock.patch("nt._supports_virtual_terminal", return_value=False),
-              unittest.mock.patch("os.isatty") as isatty_mock,
+        with (unittest.mock.patch("os.isatty") as isatty_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("_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):
+
             isatty_mock.return_value = True
             stderr_mock.fileno.return_value = 2
             stderr_mock.isatty.return_value = True
@@ -92,12 +50,18 @@ def test_colorized_detection_checks_for_environment_variables_no_vt(self):
             with unittest.mock.patch("os.environ",
                                      {'FORCE_COLOR': '1', "PYTHON_COLORS": '0'}):
                 self.assertEqual(_colorize.can_colorize(), False)
-            with unittest.mock.patch("os.environ", {}):
-                self.assertEqual(_colorize.can_colorize(), False)
 
-            isatty_mock.return_value = False
-            stderr_mock.isatty.return_value = 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)
+
+                isatty_mock.return_value = False
+                stderr_mock.isatty.return_value = False
                 self.assertEqual(_colorize.can_colorize(), False)
 
 
