Lib/test/test_winreg.py

   1 # Test the windows specific win32reg module.
   2 # Only win32reg functions not hit here: FlushKey, LoadKey and SaveKey
   3
   4 import os, sys
   5 import unittest
   6 from test import test_support
   7 threading = test_support.import_module("threading")
   8 from platform import machine
   9
  10 # Do this first so test will be skipped if module doesn't exist
  11 test_support.import_module('_winreg')
  12 # Now import everything
  13 from _winreg import *
  14
  15 try:
  16     REMOTE_NAME = sys.argv[sys.argv.index("--remote")+1]
  17 except (IndexError, ValueError):
  18     REMOTE_NAME = None
  19
  20 # tuple of (major, minor)
  21 WIN_VER = sys.getwindowsversion()[:2]
  22 # Some tests should only run on 64-bit architectures where WOW64 will be.
  23 WIN64_MACHINE = True if machine() == "AMD64" else False
  24
  25 # Starting with Windows 7 and Windows Server 2008 R2, WOW64 no longer uses
  26 # registry reflection and formerly reflected keys are shared instead.
  27 # Windows 7 and Windows Server 2008 R2 are version 6.1. Due to this, some
  28 # tests are only valid up until 6.1
  29 HAS_REFLECTION = True if WIN_VER < (6, 1) else False
  30
  31 test_key_name = "SOFTWARE\\Python Registry Test Key - Delete Me"
  32 # On OS'es that support reflection we should test with a reflected key
  33 test_reflect_key_name = "SOFTWARE\\Classes\\Python Test Key - Delete Me"
  34
  35 test_data = [
  36     ("Int Value",     45,                                      REG_DWORD),
  37     ("String Val",    "A string value",                        REG_SZ),
  38     ("StringExpand",  "The path is %path%",                    REG_EXPAND_SZ),
  39     ("Multi-string",  ["Lots", "of", "string", "values"],      REG_MULTI_SZ),
  40     ("Raw Data",      ("binary"+chr(0)+"data"),                REG_BINARY),
  41     ("Big String",    "x"*(2**14-1),                           REG_SZ),
  42     ("Big Binary",    "x"*(2**14),                             REG_BINARY),
  43 ]
  44
  45 if test_support.have_unicode:
  46     test_data += [
  47         (unicode("Unicode Val"),  unicode("A Unicode value"), REG_SZ,),
  48         ("UnicodeExpand", unicode("The path is %path%"), REG_EXPAND_SZ),
  49         ("Multi-unicode", [unicode("Lots"), unicode("of"), unicode("unicode"),
  50                            unicode("values")], REG_MULTI_SZ),
  51         ("Multi-mixed",   [unicode("Unicode"), unicode("and"), "string",
  52                            "values"], REG_MULTI_SZ),
  53     ]
  54
  55 class BaseWinregTests(unittest.TestCase):
  56
  57     def setUp(self):
  58         # Make sure that the test key is absent when the test
  59         # starts.
  60         self.delete_tree(HKEY_CURRENT_USER, test_key_name)
  61
  62     def delete_tree(self, root, subkey):
  63         try:
  64             hkey = OpenKey(root, subkey, KEY_ALL_ACCESS)
  65         except WindowsError:
  66             # subkey does not exist
  67             return
  68         while True:
  69             try:
  70                 subsubkey = EnumKey(hkey, 0)
  71             except WindowsError:
  72                 # no more subkeys
  73                 break
  74             self.delete_tree(hkey, subsubkey)
  75         CloseKey(hkey)
  76         DeleteKey(root, subkey)
  77
  78     def _write_test_data(self, root_key, CreateKey=CreateKey):
  79         # Set the default value for this key.
  80         SetValue(root_key, test_key_name, REG_SZ, "Default value")
  81         key = CreateKey(root_key, test_key_name)
  82         # Create a sub-key
  83         sub_key = CreateKey(key, "sub_key")
  84         # Give the sub-key some named values
  85
  86         for value_name, value_data, value_type in test_data:
  87             SetValueEx(sub_key, value_name, 0, value_type, value_data)
  88
  89         # Check we wrote as many items as we thought.
  90         nkeys, nvalues, since_mod = QueryInfoKey(key)
  91         self.assertEquals(nkeys, 1, "Not the correct number of sub keys")
  92         self.assertEquals(nvalues, 1, "Not the correct number of values")
  93         nkeys, nvalues, since_mod = QueryInfoKey(sub_key)
  94         self.assertEquals(nkeys, 0, "Not the correct number of sub keys")
  95         self.assertEquals(nvalues, len(test_data),
  96                           "Not the correct number of values")
  97         # Close this key this way...
  98         # (but before we do, copy the key as an integer - this allows
  99         # us to test that the key really gets closed).
 100         int_sub_key = int(sub_key)
 101         CloseKey(sub_key)
 102         try:
 103             QueryInfoKey(int_sub_key)
 104             self.fail("It appears the CloseKey() function does "
 105                       "not close the actual key!")
 106         except EnvironmentError:
 107             pass
 108         # ... and close that key that way :-)
 109         int_key = int(key)
 110         key.Close()
 111         try:
 112             QueryInfoKey(int_key)
 113             self.fail("It appears the key.Close() function "
 114                       "does not close the actual key!")
 115         except EnvironmentError:
 116             pass
 117
 118     def _read_test_data(self, root_key, OpenKey=OpenKey):
 119         # Check we can get default value for this key.
 120         val = QueryValue(root_key, test_key_name)
 121         self.assertEquals(val, "Default value",
 122                           "Registry didn't give back the correct value")
 123
 124         key = OpenKey(root_key, test_key_name)
 125         # Read the sub-keys
 126         with OpenKey(key, "sub_key") as sub_key:
 127             # Check I can enumerate over the values.
 128             index = 0
 129             while 1:
 130                 try:
 131                     data = EnumValue(sub_key, index)
 132                 except EnvironmentError:
 133                     break
 134                 self.assertIn(data, test_data,
 135                               "Didn't read back the correct test data")
 136                 index = index + 1
 137             self.assertEquals(index, len(test_data),
 138                               "Didn't read the correct number of items")
 139             # Check I can directly access each item
 140             for value_name, value_data, value_type in test_data:
 141                 read_val, read_typ = QueryValueEx(sub_key, value_name)
 142                 self.assertEquals(read_val, value_data,
 143                                   "Could not directly read the value")
 144                 self.assertEquals(read_typ, value_type,
 145                                   "Could not directly read the value")
 146         sub_key.Close()
 147         # Enumerate our main key.
 148         read_val = EnumKey(key, 0)
 149         self.assertEquals(read_val, "sub_key", "Read subkey value wrong")
 150         try:
 151             EnumKey(key, 1)
 152             self.fail("Was able to get a second key when I only have one!")
 153         except EnvironmentError:
 154             pass
 155
 156         key.Close()
 157
 158     def _delete_test_data(self, root_key):
 159         key = OpenKey(root_key, test_key_name, 0, KEY_ALL_ACCESS)
 160         sub_key = OpenKey(key, "sub_key", 0, KEY_ALL_ACCESS)
 161         # It is not necessary to delete the values before deleting
 162         # the key (although subkeys must not exist).  We delete them
 163         # manually just to prove we can :-)
 164         for value_name, value_data, value_type in test_data:
 165             DeleteValue(sub_key, value_name)
 166
 167         nkeys, nvalues, since_mod = QueryInfoKey(sub_key)
 168         self.assertEquals(nkeys, 0, "subkey not empty before delete")
 169         self.assertEquals(nvalues, 0, "subkey not empty before delete")
 170         sub_key.Close()
 171         DeleteKey(key, "sub_key")
 172
 173         try:
 174             # Shouldnt be able to delete it twice!
 175             DeleteKey(key, "sub_key")
 176             self.fail("Deleting the key twice succeeded")
 177         except EnvironmentError:
 178             pass
 179         key.Close()
 180         DeleteKey(root_key, test_key_name)
 181         # Opening should now fail!
 182         try:
 183             key = OpenKey(root_key, test_key_name)
 184             self.fail("Could open the non-existent key")
 185         except WindowsError: # Use this error name this time
 186             pass
 187
 188     def _test_all(self, root_key):
 189         self._write_test_data(root_key)
 190         self._read_test_data(root_key)
 191         self._delete_test_data(root_key)
 192
 193 class LocalWinregTests(BaseWinregTests):
 194
 195     def test_registry_works(self):
 196         self._test_all(HKEY_CURRENT_USER)
 197
 198     def test_registry_works_extended_functions(self):
 199         # Substitute the regular CreateKey and OpenKey calls with their
 200         # extended counterparts.
 201         # Note: DeleteKeyEx is not used here because it is platform dependent
 202         cke = lambda key, sub_key: CreateKeyEx(key, sub_key, 0, KEY_ALL_ACCESS)
 203         self._write_test_data(HKEY_CURRENT_USER, cke)
 204
 205         oke = lambda key, sub_key: OpenKeyEx(key, sub_key, 0, KEY_READ)
 206         self._read_test_data(HKEY_CURRENT_USER, oke)
 207
 208         self._delete_test_data(HKEY_CURRENT_USER)
 209
 210     def test_connect_registry_to_local_machine_works(self):
 211         # perform minimal ConnectRegistry test which just invokes it
 212         h = ConnectRegistry(None, HKEY_LOCAL_MACHINE)
 213         self.assertNotEqual(h.handle, 0)
 214         h.Close()
 215         self.assertEqual(h.handle, 0)
 216
 217     def test_inexistant_remote_registry(self):
 218         connect = lambda: ConnectRegistry("abcdefghijkl", HKEY_CURRENT_USER)
 219         self.assertRaises(WindowsError, connect)
 220
 221     def test_expand_environment_strings(self):
 222         r = ExpandEnvironmentStrings(u"%windir%\\test")
 223         self.assertEqual(type(r), unicode)
 224         self.assertEqual(r, os.environ["windir"] + "\\test")
 225
 226     def test_context_manager(self):
 227         # ensure that the handle is closed if an exception occurs
 228         try:
 229             with ConnectRegistry(None, HKEY_LOCAL_MACHINE) as h:
 230                 self.assertNotEqual(h.handle, 0)
 231                 raise WindowsError
 232         except WindowsError:
 233             self.assertEqual(h.handle, 0)
 234
 235     def test_changing_value(self):
 236         # Issue2810: A race condition in 2.6 and 3.1 may cause
 237         # EnumValue or QueryValue to throw "WindowsError: More data is
 238         # available"
 239         done = False
 240
 241         class VeryActiveThread(threading.Thread):
 242             def run(self):
 243                 with CreateKey(HKEY_CURRENT_USER, test_key_name) as key:
 244                     use_short = True
 245                     long_string = 'x'*2000
 246                     while not done:
 247                         s = 'x' if use_short else long_string
 248                         use_short = not use_short
 249                         SetValue(key, 'changing_value', REG_SZ, s)
 250
 251         thread = VeryActiveThread()
 252         thread.start()
 253         try:
 254             with CreateKey(HKEY_CURRENT_USER,
 255                            test_key_name+'\\changing_value') as key:
 256                 for _ in range(1000):
 257                     num_subkeys, num_values, t = QueryInfoKey(key)
 258                     for i in range(num_values):
 259                         name = EnumValue(key, i)
 260                         QueryValue(key, name[0])
 261         finally:
 262             done = True
 263             thread.join()
 264             DeleteKey(HKEY_CURRENT_USER, test_key_name+'\\changing_value')
 265             DeleteKey(HKEY_CURRENT_USER, test_key_name)
 266
 267     def test_long_key(self):
 268         # Issue2810, in 2.6 and 3.1 when the key name was exactly 256
 269         # characters, EnumKey threw "WindowsError: More data is
 270         # available"
 271         name = 'x'*256
 272         try:
 273             with CreateKey(HKEY_CURRENT_USER, test_key_name) as key:
 274                 SetValue(key, name, REG_SZ, 'x')
 275                 num_subkeys, num_values, t = QueryInfoKey(key)
 276                 EnumKey(key, 0)
 277         finally:
 278             DeleteKey(HKEY_CURRENT_USER, '\\'.join((test_key_name, name)))
 279             DeleteKey(HKEY_CURRENT_USER, test_key_name)
 280
 281     def test_dynamic_key(self):
 282         # Issue2810, when the value is dynamically generated, these
 283         # throw "WindowsError: More data is available" in 2.6 and 3.1
 284         EnumValue(HKEY_PERFORMANCE_DATA, 0)
 285         QueryValueEx(HKEY_PERFORMANCE_DATA, None)
 286
 287     # Reflection requires XP x64/Vista at a minimum. XP doesn't have this stuff
 288     # or DeleteKeyEx so make sure their use raises NotImplementedError
 289     @unittest.skipUnless(WIN_VER < (5, 2), "Requires Windows XP")
 290     def test_reflection_unsupported(self):
 291         try:
 292             with CreateKey(HKEY_CURRENT_USER, test_key_name) as ck:
 293                 self.assertNotEqual(ck.handle, 0)
 294
 295             key = OpenKey(HKEY_CURRENT_USER, test_key_name)
 296             self.assertNotEqual(key.handle, 0)
 297
 298             with self.assertRaises(NotImplementedError):
 299                 DisableReflectionKey(key)
 300             with self.assertRaises(NotImplementedError):
 301                 EnableReflectionKey(key)
 302             with self.assertRaises(NotImplementedError):
 303                 QueryReflectionKey(key)
 304             with self.assertRaises(NotImplementedError):
 305                 DeleteKeyEx(HKEY_CURRENT_USER, test_key_name)
 306         finally:
 307             DeleteKey(HKEY_CURRENT_USER, test_key_name)
 308
 309
 310 @unittest.skipUnless(REMOTE_NAME, "Skipping remote registry tests")
 311 class RemoteWinregTests(BaseWinregTests):
 312
 313     def test_remote_registry_works(self):
 314         remote_key = ConnectRegistry(REMOTE_NAME, HKEY_CURRENT_USER)
 315         self._test_all(remote_key)
 316
 317
 318 @unittest.skipUnless(WIN64_MACHINE, "x64 specific registry tests")
 319 class Win64WinregTests(BaseWinregTests):
 320
 321     def test_reflection_functions(self):
 322         # Test that we can call the query, enable, and disable functions
 323         # on a key which isn't on the reflection list with no consequences.
 324         with OpenKey(HKEY_LOCAL_MACHINE, "Software") as key:
 325             # HKLM\Software is redirected but not reflected in all OSes
 326             self.assertTrue(QueryReflectionKey(key))
 327             self.assertEquals(None, EnableReflectionKey(key))
 328             self.assertEquals(None, DisableReflectionKey(key))
 329             self.assertTrue(QueryReflectionKey(key))
 330
 331     @unittest.skipUnless(HAS_REFLECTION, "OS doesn't support reflection")
 332     def test_reflection(self):
 333         # Test that we can create, open, and delete keys in the 32-bit
 334         # area. Because we are doing this in a key which gets reflected,
 335         # test the differences of 32 and 64-bit keys before and after the
 336         # reflection occurs (ie. when the created key is closed).
 337         try:
 338             with CreateKeyEx(HKEY_CURRENT_USER, test_reflect_key_name, 0,
 339                              KEY_ALL_ACCESS | KEY_WOW64_32KEY) as created_key:
 340                 self.assertNotEqual(created_key.handle, 0)
 341
 342                 # The key should now be available in the 32-bit area
 343                 with OpenKey(HKEY_CURRENT_USER, test_reflect_key_name, 0,
 344                              KEY_ALL_ACCESS | KEY_WOW64_32KEY) as key:
 345                     self.assertNotEqual(key.handle, 0)
 346
 347                 # Write a value to what currently is only in the 32-bit area
 348                 SetValueEx(created_key, "", 0, REG_SZ, "32KEY")
 349
 350                 # The key is not reflected until created_key is closed.
 351                 # The 64-bit version of the key should not be available yet.
 352                 open_fail = lambda: OpenKey(HKEY_CURRENT_USER,
 353                                             test_reflect_key_name, 0,
 354                                             KEY_READ | KEY_WOW64_64KEY)
 355                 self.assertRaises(WindowsError, open_fail)
 356
 357             # Now explicitly open the 64-bit version of the key
 358             with OpenKey(HKEY_CURRENT_USER, test_reflect_key_name, 0,
 359                          KEY_ALL_ACCESS | KEY_WOW64_64KEY) as key:
 360                 self.assertNotEqual(key.handle, 0)
 361                 # Make sure the original value we set is there
 362                 self.assertEqual("32KEY", QueryValue(key, ""))
 363                 # Set a new value, which will get reflected to 32-bit
 364                 SetValueEx(key, "", 0, REG_SZ, "64KEY")
 365
 366             # Reflection uses a "last-writer wins policy, so the value we set
 367             # on the 64-bit key should be the same on 32-bit
 368             with OpenKey(HKEY_CURRENT_USER, test_reflect_key_name, 0,
 369                          KEY_READ | KEY_WOW64_32KEY) as key:
 370                 self.assertEqual("64KEY", QueryValue(key, ""))
 371         finally:
 372             DeleteKeyEx(HKEY_CURRENT_USER, test_reflect_key_name,
 373                         KEY_WOW64_32KEY, 0)
 374
 375     @unittest.skipUnless(HAS_REFLECTION, "OS doesn't support reflection")
 376     def test_disable_reflection(self):
 377         # Make use of a key which gets redirected and reflected
 378         try:
 379             with CreateKeyEx(HKEY_CURRENT_USER, test_reflect_key_name, 0,
 380                              KEY_ALL_ACCESS | KEY_WOW64_32KEY) as created_key:
 381                 # QueryReflectionKey returns whether or not the key is disabled
 382                 disabled = QueryReflectionKey(created_key)
 383                 self.assertEqual(type(disabled), bool)
 384                 # HKCU\Software\Classes is reflected by default
 385                 self.assertFalse(disabled)
 386
 387                 DisableReflectionKey(created_key)
 388                 self.assertTrue(QueryReflectionKey(created_key))
 389
 390             # The key is now closed and would normally be reflected to the
 391             # 64-bit area, but let's make sure that didn't happen.
 392             open_fail = lambda: OpenKeyEx(HKEY_CURRENT_USER,
 393                                           test_reflect_key_name, 0,
 394                                           KEY_READ | KEY_WOW64_64KEY)
 395             self.assertRaises(WindowsError, open_fail)
 396
 397             # Make sure the 32-bit key is actually there
 398             with OpenKeyEx(HKEY_CURRENT_USER, test_reflect_key_name, 0,
 399                            KEY_READ | KEY_WOW64_32KEY) as key:
 400                 self.assertNotEqual(key.handle, 0)
 401         finally:
 402             DeleteKeyEx(HKEY_CURRENT_USER, test_reflect_key_name,
 403                         KEY_WOW64_32KEY, 0)
 404
 405
 406 def test_main():
 407     test_support.run_unittest(LocalWinregTests, RemoteWinregTests,
 408                               Win64WinregTests)
 409
 410 if __name__ == "__main__":
 411     if not REMOTE_NAME:
 412         print "Remote registry calls can be tested using",
 413         print "'test_winreg.py --remote \\\\machine_name'"
 414     test_main()