Fix os_helper.subst_drive()

Reuse os_helper.handle_count() in test_os.test_windows.

Author: vstinner

Lib/test/support/os_helper.py

@@ -797,51 +797,61 @@ def __exit__(self, *ignore_exc):
 
 try:
     if support.MS_WINDOWS:
-        import ctypes
+        import ctypes.util
         kernel32 = ctypes.WinDLL('kernel32', use_last_error=True)
-
-        ERROR_FILE_NOT_FOUND = 2
-        DDD_REMOVE_DEFINITION = 2
-        DDD_EXACT_MATCH_ON_REMOVE = 4
-        DDD_NO_BROADCAST_SYSTEM = 8
     else:
         raise AttributeError
 except (ImportError, AttributeError):
     def subst_drive(path):
         raise unittest.SkipTest('ctypes or kernel32 is not available')
+
+    def handle_count():
+        return 0
 else:
+    ERROR_FILE_NOT_FOUND = 2
+    DDD_REMOVE_DEFINITION = 2
+    DDD_EXACT_MATCH_ON_REMOVE = 4
+    DDD_NO_BROADCAST_SYSTEM = 8
+
+
+    @ctypes.util.wrap_dll_function(kernel32)
+    def DefineDosDeviceW(
+        dwFlags: ctypes.wintypes.DWORD,
+        lpDeviceName: ctypes.c_wchar_p,
+        lpTargetPath: ctypes.c_wchar_p,
+    ) -> ctypes.wintypes.BOOL:
+        pass
+
+    @ctypes.util.wrap_dll_function(kernel32)
+    def QueryDosDeviceW(
+        lpDeviceName: ctypes.c_wchar_p,
+        lpTargetPath: ctypes.c_wchar_p,
+        ucchMax: ctypes.wintypes.DWORD,
+    ) -> ctypes.wintypes.DWORD:
+        pass
+
     @contextlib.contextmanager
     def subst_drive(path):
         """Temporarily yield a substitute drive for a given path."""
         for c in reversed(string.ascii_uppercase):
             drive = f'{c}:'
-            if (not kernel32.QueryDosDeviceW(drive, None, 0) and
+            if (not QueryDosDeviceW(drive, None, 0) and
                     ctypes.get_last_error() == ERROR_FILE_NOT_FOUND):
                 break
         else:
             raise unittest.SkipTest('no available logical drive')
-        if not kernel32.DefineDosDeviceW(
-                DDD_NO_BROADCAST_SYSTEM, drive, path):
+
+        if not DefineDosDeviceW(DDD_NO_BROADCAST_SYSTEM, drive, path):
             raise ctypes.WinError(ctypes.get_last_error())
+
         try:
             yield drive
         finally:
-            if not kernel32.DefineDosDeviceW(
-                    DDD_REMOVE_DEFINITION | DDD_EXACT_MATCH_ON_REMOVE,
-                    drive, path):
+            flags = DDD_REMOVE_DEFINITION | DDD_EXACT_MATCH_ON_REMOVE
+            if not DefineDosDeviceW(flags, drive, path):
                 raise ctypes.WinError(ctypes.get_last_error())
 
 
-try:
-    if support.MS_WINDOWS:
-        import ctypes.util
-        kernel32 = ctypes.WinDLL('kernel32', use_last_error=True)
-    else:
-        raise AttributeError
-except (ImportError, AttributeError):
-    def handle_count():
-        return 0
-else:
     @ctypes.util.wrap_dll_function(kernel32)
     def GetCurrentProcess() -> ctypes.wintypes.HANDLE:
         pass
@@ -859,6 +869,6 @@ def handle_count():
 
         handle_count = ctypes.wintypes.DWORD()
         if not GetProcessHandleCount(hproc, ctypes.byref(handle_count)):
-            raise ctypes.WinError()
+            raise ctypes.WinError(ctypes.get_last_error())
 
         return handle_count.value

Lib/test/test_os/test_windows.py

@@ -453,25 +453,8 @@ def test_unlink_removes_junction(self):
 class Win32NtTests(unittest.TestCase):
     def test_getfinalpathname_handles(self):
         nt = import_helper.import_module('nt')
-        ctypes = import_helper.import_module('ctypes')
-        # Ruff false positive -- it thinks we're redefining `ctypes` here
-        import ctypes.wintypes  # noqa: F811
 
-        kernel = ctypes.WinDLL('Kernel32.dll', use_last_error=True)
-        kernel.GetCurrentProcess.restype = ctypes.wintypes.HANDLE
-
-        kernel.GetProcessHandleCount.restype = ctypes.wintypes.BOOL
-        kernel.GetProcessHandleCount.argtypes = (ctypes.wintypes.HANDLE,
-                                                 ctypes.wintypes.LPDWORD)
-
-        # This is a pseudo-handle that doesn't need to be closed
-        hproc = kernel.GetCurrentProcess()
-
-        handle_count = ctypes.wintypes.DWORD()
-        ok = kernel.GetProcessHandleCount(hproc, ctypes.byref(handle_count))
-        self.assertEqual(1, ok)
-
-        before_count = handle_count.value
+        before_count = os_helper.handle_count()
 
         # The first two test the error path, __file__ tests the success path
         filenames = [
@@ -493,10 +476,7 @@ def test_getfinalpathname_handles(self):
                 except Exception:
                     pass
 
-        ok = kernel.GetProcessHandleCount(hproc, ctypes.byref(handle_count))
-        self.assertEqual(1, ok)
-
-        handle_delta = handle_count.value - before_count
+        handle_delta = os_helper.handle_count() - before_count
 
         self.assertEqual(0, handle_delta)