• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
1# Test the windows specific win32reg module.
2# Only win32reg functions not hit here: FlushKey, LoadKey and SaveKey
3
4import os, sys
5import unittest
6from test import test_support
7threading = test_support.import_module("threading")
8from platform import machine
9
10# Do this first so test will be skipped if module doesn't exist
11test_support.import_module('_winreg')
12# Now import everything
13from _winreg import *
14
15try:
16    REMOTE_NAME = sys.argv[sys.argv.index("--remote")+1]
17except (IndexError, ValueError):
18    REMOTE_NAME = None
19
20# tuple of (major, minor)
21WIN_VER = sys.getwindowsversion()[:2]
22# Some tests should only run on 64-bit architectures where WOW64 will be.
23WIN64_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
29HAS_REFLECTION = True if WIN_VER < (6, 1) else False
30
31test_key_name = "SOFTWARE\\Python Registry Test Key - Delete Me"
32# On OS'es that support reflection we should test with a reflected key
33test_reflect_key_name = "SOFTWARE\\Classes\\Python Test Key - Delete Me"
34
35test_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
45if 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
55class 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.assertEqual(nkeys, 1, "Not the correct number of sub keys")
92        self.assertEqual(nvalues, 1, "Not the correct number of values")
93        nkeys, nvalues, since_mod = QueryInfoKey(sub_key)
94        self.assertEqual(nkeys, 0, "Not the correct number of sub keys")
95        self.assertEqual(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.assertEqual(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.assertEqual(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.assertEqual(read_val, value_data,
143                                 "Could not directly read the value")
144                self.assertEqual(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.assertEqual(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.assertEqual(nkeys, 0, "subkey not empty before delete")
169        self.assertEqual(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
193class 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            with OpenKey(HKEY_CURRENT_USER, test_key_name, 0, KEY_ALL_ACCESS) as key:
265                DeleteKey(key, 'changing_value')
266            DeleteKey(HKEY_CURRENT_USER, test_key_name)
267
268    def test_long_key(self):
269        # Issue2810, in 2.6 and 3.1 when the key name was exactly 256
270        # characters, EnumKey threw "WindowsError: More data is
271        # available"
272        name = 'x'*256
273        try:
274            with CreateKey(HKEY_CURRENT_USER, test_key_name) as key:
275                SetValue(key, name, REG_SZ, 'x')
276                num_subkeys, num_values, t = QueryInfoKey(key)
277                EnumKey(key, 0)
278        finally:
279            with OpenKey(HKEY_CURRENT_USER, test_key_name, 0, KEY_ALL_ACCESS) as key:
280                DeleteKey(key, name)
281            DeleteKey(HKEY_CURRENT_USER, test_key_name)
282
283    def test_dynamic_key(self):
284        # Issue2810, when the value is dynamically generated, these
285        # throw "WindowsError: More data is available" in 2.6 and 3.1
286        EnumValue(HKEY_PERFORMANCE_DATA, 0)
287        QueryValueEx(HKEY_PERFORMANCE_DATA, None)
288
289    # Reflection requires XP x64/Vista at a minimum. XP doesn't have this stuff
290    # or DeleteKeyEx so make sure their use raises NotImplementedError
291    @unittest.skipUnless(WIN_VER < (5, 2), "Requires Windows XP")
292    def test_reflection_unsupported(self):
293        try:
294            with CreateKey(HKEY_CURRENT_USER, test_key_name) as ck:
295                self.assertNotEqual(ck.handle, 0)
296
297            key = OpenKey(HKEY_CURRENT_USER, test_key_name)
298            self.assertNotEqual(key.handle, 0)
299
300            with self.assertRaises(NotImplementedError):
301                DisableReflectionKey(key)
302            with self.assertRaises(NotImplementedError):
303                EnableReflectionKey(key)
304            with self.assertRaises(NotImplementedError):
305                QueryReflectionKey(key)
306            with self.assertRaises(NotImplementedError):
307                DeleteKeyEx(HKEY_CURRENT_USER, test_key_name)
308        finally:
309            DeleteKey(HKEY_CURRENT_USER, test_key_name)
310
311
312@unittest.skipUnless(REMOTE_NAME, "Skipping remote registry tests")
313class RemoteWinregTests(BaseWinregTests):
314
315    def test_remote_registry_works(self):
316        remote_key = ConnectRegistry(REMOTE_NAME, HKEY_CURRENT_USER)
317        self._test_all(remote_key)
318
319
320@unittest.skipUnless(WIN64_MACHINE, "x64 specific registry tests")
321class Win64WinregTests(BaseWinregTests):
322
323    def test_reflection_functions(self):
324        # Test that we can call the query, enable, and disable functions
325        # on a key which isn't on the reflection list with no consequences.
326        with OpenKey(HKEY_LOCAL_MACHINE, "Software") as key:
327            # HKLM\Software is redirected but not reflected in all OSes
328            self.assertTrue(QueryReflectionKey(key))
329            self.assertEqual(None, EnableReflectionKey(key))
330            self.assertEqual(None, DisableReflectionKey(key))
331            self.assertTrue(QueryReflectionKey(key))
332
333    @unittest.skipUnless(HAS_REFLECTION, "OS doesn't support reflection")
334    def test_reflection(self):
335        # Test that we can create, open, and delete keys in the 32-bit
336        # area. Because we are doing this in a key which gets reflected,
337        # test the differences of 32 and 64-bit keys before and after the
338        # reflection occurs (ie. when the created key is closed).
339        try:
340            with CreateKeyEx(HKEY_CURRENT_USER, test_reflect_key_name, 0,
341                             KEY_ALL_ACCESS | KEY_WOW64_32KEY) as created_key:
342                self.assertNotEqual(created_key.handle, 0)
343
344                # The key should now be available in the 32-bit area
345                with OpenKey(HKEY_CURRENT_USER, test_reflect_key_name, 0,
346                             KEY_ALL_ACCESS | KEY_WOW64_32KEY) as key:
347                    self.assertNotEqual(key.handle, 0)
348
349                # Write a value to what currently is only in the 32-bit area
350                SetValueEx(created_key, "", 0, REG_SZ, "32KEY")
351
352                # The key is not reflected until created_key is closed.
353                # The 64-bit version of the key should not be available yet.
354                open_fail = lambda: OpenKey(HKEY_CURRENT_USER,
355                                            test_reflect_key_name, 0,
356                                            KEY_READ | KEY_WOW64_64KEY)
357                self.assertRaises(WindowsError, open_fail)
358
359            # Now explicitly open the 64-bit version of the key
360            with OpenKey(HKEY_CURRENT_USER, test_reflect_key_name, 0,
361                         KEY_ALL_ACCESS | KEY_WOW64_64KEY) as key:
362                self.assertNotEqual(key.handle, 0)
363                # Make sure the original value we set is there
364                self.assertEqual("32KEY", QueryValue(key, ""))
365                # Set a new value, which will get reflected to 32-bit
366                SetValueEx(key, "", 0, REG_SZ, "64KEY")
367
368            # Reflection uses a "last-writer wins policy, so the value we set
369            # on the 64-bit key should be the same on 32-bit
370            with OpenKey(HKEY_CURRENT_USER, test_reflect_key_name, 0,
371                         KEY_READ | KEY_WOW64_32KEY) as key:
372                self.assertEqual("64KEY", QueryValue(key, ""))
373        finally:
374            DeleteKeyEx(HKEY_CURRENT_USER, test_reflect_key_name,
375                        KEY_WOW64_32KEY, 0)
376
377    @unittest.skipUnless(HAS_REFLECTION, "OS doesn't support reflection")
378    def test_disable_reflection(self):
379        # Make use of a key which gets redirected and reflected
380        try:
381            with CreateKeyEx(HKEY_CURRENT_USER, test_reflect_key_name, 0,
382                             KEY_ALL_ACCESS | KEY_WOW64_32KEY) as created_key:
383                # QueryReflectionKey returns whether or not the key is disabled
384                disabled = QueryReflectionKey(created_key)
385                self.assertEqual(type(disabled), bool)
386                # HKCU\Software\Classes is reflected by default
387                self.assertFalse(disabled)
388
389                DisableReflectionKey(created_key)
390                self.assertTrue(QueryReflectionKey(created_key))
391
392            # The key is now closed and would normally be reflected to the
393            # 64-bit area, but let's make sure that didn't happen.
394            open_fail = lambda: OpenKeyEx(HKEY_CURRENT_USER,
395                                          test_reflect_key_name, 0,
396                                          KEY_READ | KEY_WOW64_64KEY)
397            self.assertRaises(WindowsError, open_fail)
398
399            # Make sure the 32-bit key is actually there
400            with OpenKeyEx(HKEY_CURRENT_USER, test_reflect_key_name, 0,
401                           KEY_READ | KEY_WOW64_32KEY) as key:
402                self.assertNotEqual(key.handle, 0)
403        finally:
404            DeleteKeyEx(HKEY_CURRENT_USER, test_reflect_key_name,
405                        KEY_WOW64_32KEY, 0)
406
407
408def test_main():
409    test_support.run_unittest(LocalWinregTests, RemoteWinregTests,
410                              Win64WinregTests)
411
412if __name__ == "__main__":
413    if not REMOTE_NAME:
414        print "Remote registry calls can be tested using",
415        print "'test_winreg.py --remote \\\\machine_name'"
416    test_main()
417