# Test the windows specific win32reg module.
# Only win32reg functions not hit here: FlushKey, LoadKey and SaveKey

from _winreg import *
import os, sys

test_key_name = "SOFTWARE\\Python Registry Test Key - Delete Me"

test_data = [
    ("Int Value",     45,                                      REG_DWORD),
    ("String Val",    "A string value",                        REG_SZ,),
    (u"Unicode Val",  u"A Unicode value",                      REG_SZ,),
    ("StringExpand",  "The path is %path%",                    REG_EXPAND_SZ),
    ("UnicodeExpand", u"The path is %path%",                   REG_EXPAND_SZ),
    ("Multi-string",  ["Lots", "of", "string", "values"],      REG_MULTI_SZ),
    ("Multi-unicode", [u"Lots", u"of", u"unicode", u"values"], REG_MULTI_SZ),
    ("Multi-mixed",   [u"Unicode", u"and", "string", "values"],REG_MULTI_SZ),
    ("Raw Data",      ("binary"+chr(0)+"data"),                REG_BINARY),
]

def WriteTestData(root_key):
    # Set the default value for this key.
    SetValue(root_key, test_key_name, REG_SZ, "Default value")
    key = CreateKey(root_key, test_key_name)
    # Create a sub-key
    sub_key = CreateKey(key, "sub_key")
    # Give the sub-key some named values

    for value_name, value_data, value_type in test_data:
        SetValueEx(sub_key, value_name, 0, value_type, value_data)

    # Check we wrote as many items as we thought.
    nkeys, nvalues, since_mod = QueryInfoKey(key)
    assert nkeys==1, "Not the correct number of sub keys"
    assert nvalues==1, "Not the correct number of values"
    nkeys, nvalues, since_mod = QueryInfoKey(sub_key)
    assert nkeys==0, "Not the correct number of sub keys"
    assert nvalues==len(test_data), "Not the correct number of values"
    # Close this key this way...
    # (but before we do, copy the key as an integer - this allows
    # us to test that the key really gets closed).
    int_sub_key = int(sub_key)
    CloseKey(sub_key)
    try:
        QueryInfoKey(int_sub_key)
        raise RuntimeError, "It appears the CloseKey() function does not close the actual key!"
    except EnvironmentError:
        pass
    # ... and close that key that way :-)
    int_key = int(key)
    key.Close()
    try:
        QueryInfoKey(int_key)
        raise RuntimeError, "It appears the key.Close() function does not close the actual key!"
    except EnvironmentError:
        pass

def ReadTestData(root_key):
    # Check we can get default value for this key.
    val = QueryValue(root_key, test_key_name)
    assert val=="Default value", "Registry didn't give back the correct value"

    key = OpenKey(root_key, test_key_name)
    # Read the sub-keys
    sub_key = OpenKey(key, "sub_key")
    # Check I can enumerate over the values.
    index = 0
    while 1:
        try:
            data = EnumValue(sub_key, index)
        except EnvironmentError:
            break
        assert data in test_data, "Didn't read back the correct test data"
        index = index + 1
    assert index==len(test_data), "Didn't read the correct number of items"
    # Check I can directly access each item
    for value_name, value_data, value_type in test_data:
        read_val, read_typ = QueryValueEx(sub_key, value_name)
        assert read_val==value_data and read_typ == value_type, \
               "Could not directly read the value"
    sub_key.Close()
    # Enumerate our main key.
    read_val = EnumKey(key, 0)
    assert read_val == "sub_key", "Read subkey value wrong"
    try:
        EnumKey(key, 1)
        assert 0, "Was able to get a second key when I only have one!"
    except EnvironmentError:
        pass

    key.Close()

def DeleteTestData(root_key):
    key = OpenKey(root_key, test_key_name, 0, KEY_ALL_ACCESS)
    sub_key = OpenKey(key, "sub_key", 0, KEY_ALL_ACCESS)
    # It is not necessary to delete the values before deleting
    # the key (although subkeys must not exist).  We delete them
    # manually just to prove we can :-)
    for value_name, value_data, value_type in test_data:
        DeleteValue(sub_key, value_name)

    nkeys, nvalues, since_mod = QueryInfoKey(sub_key)
    assert nkeys==0 and nvalues==0, "subkey not empty before delete"
    sub_key.Close()
    DeleteKey(key, "sub_key")

    try:
        # Shouldnt be able to delete it twice!
        DeleteKey(key, "sub_key")
        assert 0, "Deleting the key twice succeeded"
    except EnvironmentError:
        pass
    key.Close()
    DeleteKey(root_key, test_key_name)
    # Opening should now fail!
    try:
        key = OpenKey(root_key, test_key_name)
        assert 0, "Could open the non-existent key"
    except WindowsError: # Use this error name this time
        pass

def TestAll(root_key):
    WriteTestData(root_key)
    ReadTestData(root_key)
    DeleteTestData(root_key)

# Test on my local machine.
TestAll(HKEY_CURRENT_USER)
print "Local registry tests worked"
try:
    remote_name = sys.argv[sys.argv.index("--remote")+1]
except (IndexError, ValueError):
    remote_name = None

if remote_name is not None:
    try:
        remote_key = ConnectRegistry(remote_name, HKEY_CURRENT_USER)
    except EnvironmentError, exc:
        print "Could not connect to the remote machine -", exc.strerror
        remote_key = None
    if remote_key is not None:
        TestAll(remote_key)
        print "Remote registry tests worked"
else:
    print "Remote registry calls can be tested using",
    print "'test_winreg.py --remote \\\\machine_name'"