Migrate command: wallet rebuild (#747)

Signed-off-by: Guillaume George <lysandergc@gmail.com>

Author: LysanderGG

neo/Implementations/Wallets/peewee/UserWallet.py

@@ -63,9 +63,9 @@ def Migrate(self):
     def DB(self):
         return self._db
 
-    def Rebuild(self):
+    def Rebuild(self, start_block=0):
         try:
-            super(UserWallet, self).Rebuild()
+            super(UserWallet, self).Rebuild(start_block)
 
             logger.debug("wallet rebuild: deleting %s coins and %s transactions" %
                          (Coin.select().count(), Transaction.select().count()))

neo/Prompt/Commands/Wallet.py

@@ -38,6 +38,7 @@ def __init__(self):
         self.register_sub_command(CommandWalletSend())
         self.register_sub_command(CommandWalletSendMany())
         self.register_sub_command(CommandWalletSign())
+        self.register_sub_command(CommandWalletRebuild())
 
     def command_desc(self):
         return CommandDesc('wallet', 'manage wallets')
@@ -210,6 +211,28 @@ def command_desc(self):
         return CommandDesc('create_addr', 'create a new wallet address', params=[p1])
 
 
+class CommandWalletRebuild(CommandBase):
+
+    def __init__(self):
+        super().__init__()
+
+    def execute(self, arguments):
+        PromptData.Prompt.stop_wallet_loop()
+
+        start_block = get_arg(arguments, 0, convert_to_int=True)
+        if not start_block or start_block < 0:
+            start_block = 0
+        print(f"Restarting at block {start_block}")
+
+        PromptData.Wallet.Rebuild(start_block)
+
+        PromptData.Prompt.start_wallet_loop()
+
+    def command_desc(self):
+        p1 = ParameterDesc('start_block', 'block number to start the resync at', optional=True)
+        return CommandDesc('rebuild', 'rebuild the wallet index', params=[p1])
+
+
 #########################################################################
 #########################################################################
 

neo/Prompt/Commands/tests/test_wallet_commands.py

@@ -240,6 +240,26 @@ def test_wallet_create_address(self):
         self.assertEqual(type(res), UserWallet)
         self.assertEqual(len(res.Addresses), 9)
 
+    def test_wallet_rebuild(self):
+        with patch('neo.Prompt.PromptData.PromptData.Prompt'):
+            # test wallet rebuild with no wallet open
+            args = ['rebuild']
+            res = CommandWallet().execute(args)
+            self.assertFalse(res)
+
+            self.OpenWallet()
+            PromptData.Wallet._current_height = 12345
+
+            # test wallet rebuild with no argument
+            args = ['rebuild']
+            CommandWallet().execute(args)
+            self.assertEqual(PromptData.Wallet._current_height, 0)
+
+            # test wallet rebuild with start block
+            args = ['rebuild', '42']
+            CommandWallet().execute(args)
+            self.assertEqual(PromptData.Wallet._current_height, 42)
+
     ##########################################################
     ##########################################################
     def test_1_import_addr(self):

neo/Wallets/Wallet.py

@@ -749,13 +749,13 @@ def ProcessNewBlock(self, block):
             traceback.print_exc()
             logger.error("could not process %s " % e)
 
-    def Rebuild(self):
+    def Rebuild(self, start_block=0):
         """
         Sets the current height to 0 and now `ProcessBlocks` will start from
         the beginning of the blockchain.
         """
         self._coins = {}
-        self._current_height = 0
+        self._current_height = start_block
 
     def OnProcessNewBlock(self, block, added, changed, deleted):
         # abstract