gh-154137: Add test.support.os_helper.handle_count() (#157637)

Use the new @ctypes.util.wrap_dll_function decorator to wrap kernel32
functions.

Author: vstinner

Lib/test/support/os_helper.py

@@ -13,6 +13,8 @@
 import warnings
 
 from test import support
+if support.MS_WINDOWS:
+    import _winapi
 
 
 # Filename used for testing
@@ -854,36 +856,72 @@ def __new__(mcls, name, bases, dct, *, source_date_epoch):
 
 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())
+
+    @ctypes.util.wrap_dll_function(kernel32)
+    def GetProcessHandleCount(khProcess: ctypes.wintypes.HANDLE,
+                              pdwHandleCount: ctypes.wintypes.LPDWORD) -> ctypes.wintypes.BOOL:
+        pass
+
+    del kernel32
+
+    def handle_count():
+        # Pseudo-handle that doesn't need to be closed
+        hproc = _winapi.GetCurrentProcess()
+
+        handle_count = ctypes.wintypes.DWORD()
+        if not GetProcessHandleCount(hproc, ctypes.byref(handle_count)):
+            raise ctypes.WinError(ctypes.get_last_error())
+
+        return handle_count.value

Lib/test/test_os/test_windows.py

@@ -457,28 +457,7 @@ 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)
-        @ctypes.util.wrap_dll_function(kernel)
-        def GetCurrentProcess() -> ctypes.wintypes.HANDLE:
-            pass
-
-        @ctypes.util.wrap_dll_function(kernel)
-        def GetProcessHandleCount(khProcess: ctypes.wintypes.HANDLE,
-                                  pdwHandleCount: ctypes.wintypes.LPDWORD) -> ctypes.wintypes.BOOL:
-            pass
-
-        # This is a pseudo-handle that doesn't need to be closed
-        hproc = GetCurrentProcess()
-
-        handle_count = ctypes.wintypes.DWORD()
-        ok = 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 = [
@@ -500,11 +479,7 @@ def GetProcessHandleCount(khProcess: ctypes.wintypes.HANDLE,
                 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)
 
     @support.requires_subprocess()

Lib/test/test_support.py

@@ -27,6 +27,10 @@
 from test.support import socket_helper
 from test.support import warnings_helper
 
+if support.MS_WINDOWS:
+    import _winapi
+
+
 TESTFN = os_helper.TESTFN
 
 
@@ -624,6 +628,20 @@ def test_fd_count(self):
             os.close(fd)
         self.assertEqual(more - start, 1)
 
+    @unittest.skipUnless(support.MS_WINDOWS, "test specific to Windows")
+    def test_handle_count(self):
+        start = os_helper.handle_count()
+        handle = _winapi.CreateFile(
+                        __file__, _winapi.GENERIC_READ,
+                        0, _winapi.NULL,
+                        _winapi.OPEN_EXISTING,
+                        0, _winapi.NULL)
+        try:
+            more = os_helper.handle_count()
+        finally:
+            _winapi.CloseHandle(handle)
+        self.assertEqual(more - start, 1)
+
     def check_print_warning(self, msg, expected):
         stderr = io.StringIO()
         with support.swap_attr(support.print_warning, 'orig_stderr', stderr):