bpo-39007: Add auditing events to functions in winreg (GH-17541)

Also allows winreg.CloseKey() to accept same types as other functions.

2 files changed, 37 insertions, 0 deletions

diff --git a/Lib/test/audit-tests.py b/Lib/test/audit-tests.py
index ed08612..33f3209 100644
--- a/Lib/test/audit-tests.py
+++ b/Lib/test/audit-tests.py
@@ -304,6 +304,29 @@ def test_unraisablehook():
     write_unraisable_exc(RuntimeError("nonfatal-error"), "for audit hook test", None)
 
 
+def test_winreg():
+    from winreg import OpenKey, EnumKey, CloseKey, HKEY_LOCAL_MACHINE
+
+    def hook(event, args):
+        if not event.startswith("winreg."):
+            return
+        print(event, *args)
+
+    sys.addaudithook(hook)
+
+    k = OpenKey(HKEY_LOCAL_MACHINE, "Software")
+    EnumKey(k, 0)
+    try:
+        EnumKey(k, 10000)
+    except OSError:
+        pass
+    else:
+        raise RuntimeError("Expected EnumKey(HKLM, 10000) to fail")
+
+    kv = k.Detach()
+    CloseKey(kv)
+
+
 if __name__ == "__main__":
     from test.libregrtest.setup import suppress_msvcrt_asserts
 
diff --git a/Lib/test/test_audit.py b/Lib/test/test_audit.py
index 31a0855..73dd5c5 100644
--- a/Lib/test/test_audit.py
+++ b/Lib/test/test_audit.py
@@ -104,6 +104,20 @@ class AuditTest(unittest.TestCase):
             "RuntimeError('nonfatal-error') Exception ignored for audit hook test",
         )
 
+    def test_winreg(self):
+        support.import_module("winreg")
+        returncode, events, stderr = self.run_python("test_winreg")
+        if returncode:
+            self.fail(stderr)
+
+        self.assertEqual(events[0][0], "winreg.OpenKey")
+        self.assertEqual(events[1][0], "winreg.OpenKey/result")
+        expected = events[1][2]
+        self.assertTrue(expected)
+        self.assertSequenceEqual(["winreg.EnumKey", " ", f"{expected} 0"], events[2])
+        self.assertSequenceEqual(["winreg.EnumKey", " ", f"{expected} 10000"], events[3])
+        self.assertSequenceEqual(["winreg.PyHKEY.Detach", " ", expected], events[4])
+
 
 if __name__ == "__main__":
     unittest.main()