Skip to content
This repository was archived by the owner on Nov 15, 2021. It is now read-only.

[refactor-prompt] Add missing tests for wallet create, verbose and create_addr #733

Merged
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
121 changes: 59 additions & 62 deletions neo/Prompt/Commands/Wallet.py
Original file line number Diff line number Diff line change
Expand Up @@ -56,7 +56,7 @@ def execute(self, arguments):

if not item:
print("Wallet %s " % json.dumps(wallet.ToJson(), indent=4))
return
return wallet

try:
return self.execute_sub_command(item, arguments[1:])
Expand All @@ -75,45 +75,43 @@ def execute(self, arguments):
PromptData.close_wallet()
path = get_arg(arguments, 0)

if path:
if os.path.exists(path):
print("File already exists")
return

passwd1 = prompt("[password]> ", is_password=True)
passwd2 = prompt("[password again]> ", is_password=True)

if passwd1 != passwd2 or len(passwd1) < 10:
print("Please provide matching passwords that are at least 10 characters long")
return

password_key = to_aes_key(passwd1)

try:
PromptData.Wallet = UserWallet.Create(path=path, password=password_key)
contract = PromptData.Wallet.GetDefaultContract()
key = PromptData.Wallet.GetKey(contract.PublicKeyHash)
print("Wallet %s" % json.dumps(PromptData.Wallet.ToJson(), indent=4))
print("Pubkey %s" % key.PublicKey.encode_point(True))
return PromptData.Wallet
except Exception as e:
print("Exception creating wallet: %s" % e)
PromptData.Wallet = None
if os.path.isfile(path):
try:
os.remove(path)
except Exception as e:
print("Could not remove {}: {}".format(path, e))
return

if PromptData.Wallet:
PromptData.Prompt.start_wallet_loop()
return

else:
if not path:
print("Please specify a path")
return

if os.path.exists(path):
print("File already exists")
return

passwd1 = prompt("[password]> ", is_password=True)
passwd2 = prompt("[password again]> ", is_password=True)

if passwd1 != passwd2 or len(passwd1) < 10:
print("Please provide matching passwords that are at least 10 characters long")
return

password_key = to_aes_key(passwd1)

try:
PromptData.Wallet = UserWallet.Create(path=path, password=password_key)
contract = PromptData.Wallet.GetDefaultContract()
key = PromptData.Wallet.GetKey(contract.PublicKeyHash)
print("Wallet %s" % json.dumps(PromptData.Wallet.ToJson(), indent=4))
print("Pubkey %s" % key.PublicKey.encode_point(True))
except Exception as e:
print("Exception creating wallet: %s" % e)
PromptData.Wallet = None
if os.path.isfile(path):
try:
os.remove(path)
except Exception as e:
print("Could not remove {}: {}".format(path, e))
return

if PromptData.Wallet:
PromptData.Prompt.start_wallet_loop()
return PromptData.Wallet

def command_desc(self):
p1 = ParameterDesc('path', 'path to store the wallet file')
return CommandDesc('create', 'create a new NEO wallet (with 1 address)', [p1])
Expand All @@ -130,28 +128,25 @@ def execute(self, arguments):

path = get_arg(arguments, 0)

if path:

if not os.path.exists(path):
print("Wallet file not found")
return
if not path:
print("Please specify a path")
return

passwd = prompt("[password]> ", is_password=True)
password_key = to_aes_key(passwd)
if not os.path.exists(path):
print("Wallet file not found")
return

try:
PromptData.Wallet = UserWallet.Open(path, password_key)
passwd = prompt("[password]> ", is_password=True)
password_key = to_aes_key(passwd)

PromptData.Prompt.start_wallet_loop()
print("Opened wallet at %s" % path)
return PromptData.Wallet
except Exception as e:
print("Could not open wallet: %s" % e)
return
try:
PromptData.Wallet = UserWallet.Open(path, password_key)

else:
print("Please specify a path")
return
PromptData.Prompt.start_wallet_loop()
print("Opened wallet at %s" % path)
return PromptData.Wallet
except Exception as e:
print("Could not open wallet: %s" % e)

def command_desc(self):
p1 = ParameterDesc('path', 'path to open the wallet file')
Expand All @@ -175,7 +170,7 @@ class CommandWalletVerbose(CommandBase):
def __init__(self):
super().__init__()

def execute(self, arguments):
def execute(self, arguments=None):
print("Wallet %s " % json.dumps(PromptData.Wallet.ToJson(verbose=True), indent=4))
return True

Expand All @@ -188,12 +183,9 @@ class CommandWalletMigrate(CommandBase):
def __init__(self):
super().__init__()

def execute(self, arguments):
if PromptData.Wallet is not None:
PromptData.Wallet.Migrate()
print("Migrated wallet")
return True
return False
def execute(self, arguments=None):
PromptData.Wallet.Migrate()
return True

def command_desc(self):
return CommandDesc('migrate', 'migrate an old wallet to the new format')
Expand All @@ -206,6 +198,11 @@ def __init__(self):

def execute(self, arguments):
addresses_to_create = get_arg(arguments, 0)

if not addresses_to_create:
print("Please specify a number of addresses to create.")
return

return CreateAddress(PromptData.Wallet, addresses_to_create)

def command_desc(self):
Expand Down
153 changes: 150 additions & 3 deletions neo/Prompt/Commands/tests/test_wallet_commands.py
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,10 @@
from neo.Core.Blockchain import Blockchain
from neocore.UInt160 import UInt160
from neocore.Fixed8 import Fixed8
from neo.Prompt.Commands.Wallet import CommandWallet, CreateAddress, DeleteAddress, ImportToken, ImportWatchAddr, ShowUnspentCoins, SplitUnspentCoin
from neo.Prompt.Commands.Wallet import CommandWallet
from neo.Prompt.Commands.Wallet import CreateAddress, DeleteAddress, ImportToken, ImportWatchAddr, ShowUnspentCoins, SplitUnspentCoin
from neo.Prompt.PromptData import PromptData
import os
import shutil
from mock import patch

Expand Down Expand Up @@ -34,8 +37,109 @@ def GetWallet1(cls, recreate=False):
to_aes_key(UserWalletTestCase.wallet_1_pass()))
return cls._wallet1

@classmethod
def OpenWallet(cls):
PromptData.Wallet = cls.GetWallet1(recreate=True)

@classmethod
def tearDown(cls):
PromptData.Wallet = None

# Beginning with refactored tests

def test_wallet(self):
# without wallet opened
res = CommandWallet().execute(None)
self.assertFalse(res)

# with wallet opened
self.OpenWallet()
res = CommandWallet().execute(None)
self.assertEqual(type(res), UserWallet)

def test_wallet_wrong_command(self):
self.OpenWallet()
args = ['badcommand']
res = CommandWallet().execute(args)
self.assertFalse(res)

def test_wallet_create(self):
def remove_new_wallet():
path = UserWalletTestCase.new_wallet_dest()
try:
if os.path.exists(path):
os.remove(path)
except Exception as e:
print("couldn't remove wallets %s " % e)

with patch('neo.Prompt.PromptData.PromptData.Prompt'):
with patch('neo.Prompt.Commands.Wallet.prompt', side_effect=["testpassword", "testpassword"]):
# test wallet create successful
path = UserWalletTestCase.new_wallet_dest()
args = ['create', path]
self.assertFalse(os.path.isfile(path))
res = CommandWallet().execute(args)
self.assertEqual(type(res), UserWallet)
self.assertTrue(os.path.isfile(path))
remove_new_wallet()

# test wallet create with no path
args = ['create']
res = CommandWallet().execute(args)
self.assertFalse(res)

# test wallet open with already existing path
with patch('neo.Prompt.Commands.Wallet.prompt', side_effect=["testpassword", "testpassword"]):
path = UserWalletTestCase.new_wallet_dest()
args = ['create', path]
self.assertFalse(os.path.isfile(path))
res = CommandWallet().execute(args)
self.assertEqual(type(res), UserWallet)
self.assertTrue(os.path.isfile(path))

res = CommandWallet().execute(args)
self.assertFalse(res)
self.assertTrue(os.path.isfile(path))
remove_new_wallet()

# test wallet with different passwords
with patch('neo.Prompt.Commands.Wallet.prompt', side_effect=["testpassword", "bad"]):
path = UserWalletTestCase.new_wallet_dest()
args = ['create', path]
self.assertFalse(os.path.isfile(path))
res = CommandWallet().execute(args)
self.assertFalse(res)
self.assertFalse(os.path.isfile(path))

# test wallet create unsuccessful
with patch('neo.Prompt.Commands.Wallet.prompt', side_effect=["testpassword", "testpassword"]):
with patch('neo.Implementations.Wallets.peewee.UserWallet.UserWallet.Create', side_effect=[Exception('test exception')]):
path = UserWalletTestCase.new_wallet_dest()
args = ['create', path]
res = CommandWallet().execute(args)
self.assertFalse(res)
self.assertFalse(os.path.isfile(path))

# test wallet create exception after creation
with patch('neo.Prompt.Commands.Wallet.prompt', side_effect=["testpassword", "testpassword"]):
with patch('neo.Wallets.Wallet.Wallet.GetKey', side_effect=[Exception('test exception')]):
path = UserWalletTestCase.new_wallet_dest()
args = ['create', path]
res = CommandWallet().execute(args)
self.assertFalse(res)
self.assertFalse(os.path.isfile(path))

# test wallet create exception after creation with file deletion failure
with patch('neo.Prompt.Commands.Wallet.prompt', side_effect=["testpassword", "testpassword"]):
with patch('neo.Wallets.Wallet.Wallet.GetKey', side_effect=[Exception('test exception')]):
with patch('os.remove', side_effect=[Exception('test exception')]):
path = UserWalletTestCase.new_wallet_dest()
args = ['create', path]
res = CommandWallet().execute(args)
self.assertFalse(res)
self.assertTrue(os.path.isfile(path))
remove_new_wallet()

def test_wallet_open(self):
with patch('neo.Prompt.PromptData.PromptData.Prompt'):
with patch('neo.Prompt.Commands.Wallet.prompt', side_effect=["testpassword"]):
Expand All @@ -44,7 +148,7 @@ def test_wallet_open(self):

res = CommandWallet().execute(args)

self.assertEqual(str(type(res)), "<class 'neo.Implementations.Wallets.peewee.UserWallet.UserWallet'>")
self.assertEqual(type(res), UserWallet)

# test wallet open with no path; this will also close the open wallet
args = ['open']
Expand Down Expand Up @@ -84,7 +188,7 @@ def test_wallet_close(self):

res = CommandWallet().execute(args)

self.assertEqual(str(type(res)), "<class 'neo.Implementations.Wallets.peewee.UserWallet.UserWallet'>")
self.assertEqual(type(res), UserWallet)

# now close the open wallet manually
args = ['close']
Expand All @@ -93,6 +197,49 @@ def test_wallet_close(self):

self.assertTrue(res)

def test_wallet_verbose(self):
# test wallet verbose with no wallet opened
args = ['verbose']
res = CommandWallet().execute(args)
self.assertFalse(res)

# test wallet close with open wallet
self.OpenWallet()
args = ['verbose']
res = CommandWallet().execute(args)
self.assertTrue(res)

def test_wallet_create_address(self):
# test wallet create address with no wallet open
args = ['create_addr', 1]
res = CommandWallet().execute(args)
self.assertFalse(res)

self.OpenWallet()

# test wallet create address with no argument
args = ['create_addr']
res = CommandWallet().execute(args)
self.assertFalse(res)

# test wallet create address with negative number
args = ['create_addr', -1]
res = CommandWallet().execute(args)
self.assertFalse(res)

# test wallet create successful
args = ['create_addr', 1]
res = CommandWallet().execute(args)
self.assertTrue(res)
self.assertEqual(type(res), UserWallet)
self.assertEqual(len(res.Addresses), 2) # Has one address when created.

args = ['create_addr', 7]
res = CommandWallet().execute(args)
self.assertTrue(res)
self.assertEqual(type(res), UserWallet)
self.assertEqual(len(res.Addresses), 9)

##########################################################
##########################################################
def test_1_import_addr(self):
Expand Down