github.com/prysmaticlabs/prysm@v1.4.4/validator/accounts/wallet_recover_test.go (about)

     1  package accounts
     2  
     3  import (
     4  	"context"
     5  	"flag"
     6  	"io/ioutil"
     7  	"os"
     8  	"path/filepath"
     9  	"strconv"
    10  	"testing"
    11  
    12  	"github.com/prysmaticlabs/prysm/cmd/validator/flags"
    13  	"github.com/prysmaticlabs/prysm/shared/testutil/assert"
    14  	"github.com/prysmaticlabs/prysm/shared/testutil/require"
    15  	"github.com/prysmaticlabs/prysm/validator/accounts/iface"
    16  	"github.com/prysmaticlabs/prysm/validator/accounts/wallet"
    17  	"github.com/prysmaticlabs/prysm/validator/keymanager"
    18  	"github.com/prysmaticlabs/prysm/validator/keymanager/derived"
    19  	"github.com/urfave/cli/v2"
    20  )
    21  
    22  type recoverCfgStruct struct {
    23  	walletDir        string
    24  	passwordFilePath string
    25  	mnemonicFilePath string
    26  	numAccounts      int64
    27  }
    28  
    29  func setupRecoverCfg(t *testing.T) *recoverCfgStruct {
    30  	testDir := t.TempDir()
    31  	walletDir := filepath.Join(testDir, walletDirName)
    32  	passwordFilePath := filepath.Join(testDir, passwordFileName)
    33  	require.NoError(t, ioutil.WriteFile(passwordFilePath, []byte(password), os.ModePerm))
    34  	mnemonicFilePath := filepath.Join(testDir, mnemonicFileName)
    35  	require.NoError(t, ioutil.WriteFile(mnemonicFilePath, []byte(mnemonic), os.ModePerm))
    36  
    37  	return &recoverCfgStruct{
    38  		walletDir:        walletDir,
    39  		passwordFilePath: passwordFilePath,
    40  		mnemonicFilePath: mnemonicFilePath,
    41  	}
    42  }
    43  
    44  func createRecoverCliCtx(t *testing.T, cfg *recoverCfgStruct) *cli.Context {
    45  	app := cli.App{}
    46  	set := flag.NewFlagSet("test", 0)
    47  	set.String(flags.WalletDirFlag.Name, cfg.walletDir, "")
    48  	set.String(flags.WalletPasswordFileFlag.Name, cfg.passwordFilePath, "")
    49  	set.String(flags.KeymanagerKindFlag.Name, keymanager.Derived.String(), "")
    50  	set.String(flags.MnemonicFileFlag.Name, cfg.mnemonicFilePath, "")
    51  	set.Bool(flags.SkipMnemonic25thWordCheckFlag.Name, true, "")
    52  	set.Int64(flags.NumAccountsFlag.Name, cfg.numAccounts, "")
    53  	assert.NoError(t, set.Set(flags.SkipMnemonic25thWordCheckFlag.Name, "true"))
    54  	assert.NoError(t, set.Set(flags.WalletDirFlag.Name, cfg.walletDir))
    55  	assert.NoError(t, set.Set(flags.WalletPasswordFileFlag.Name, cfg.passwordFilePath))
    56  	assert.NoError(t, set.Set(flags.KeymanagerKindFlag.Name, keymanager.Derived.String()))
    57  	assert.NoError(t, set.Set(flags.MnemonicFileFlag.Name, cfg.mnemonicFilePath))
    58  	assert.NoError(t, set.Set(flags.NumAccountsFlag.Name, strconv.Itoa(int(cfg.numAccounts))))
    59  	return cli.NewContext(&app, set, nil)
    60  }
    61  
    62  func TestRecoverDerivedWallet(t *testing.T) {
    63  	cfg := setupRecoverCfg(t)
    64  	cfg.numAccounts = 4
    65  	cliCtx := createRecoverCliCtx(t, cfg)
    66  	require.NoError(t, RecoverWalletCli(cliCtx))
    67  
    68  	ctx := context.Background()
    69  	w, err := wallet.OpenWallet(cliCtx.Context, &wallet.Config{
    70  		WalletDir:      cfg.walletDir,
    71  		WalletPassword: password,
    72  	})
    73  	assert.NoError(t, err)
    74  
    75  	km, err := w.InitializeKeymanager(cliCtx.Context, iface.InitKeymanagerConfig{ListenForChanges: false})
    76  	require.NoError(t, err)
    77  	derivedKM, ok := km.(*derived.Keymanager)
    78  	if !ok {
    79  		t.Fatal("not a derived keymanager")
    80  	}
    81  	names, err := derivedKM.ValidatingAccountNames(ctx)
    82  	assert.NoError(t, err)
    83  	require.Equal(t, len(names), int(cfg.numAccounts))
    84  }
    85  
    86  // TestRecoverDerivedWallet_OneAccount is a test for regression in cases where the number of accounts recovered is 1
    87  func TestRecoverDerivedWallet_OneAccount(t *testing.T) {
    88  	cfg := setupRecoverCfg(t)
    89  	cfg.numAccounts = 1
    90  	cliCtx := createRecoverCliCtx(t, cfg)
    91  	require.NoError(t, RecoverWalletCli(cliCtx))
    92  
    93  	_, err := wallet.OpenWallet(cliCtx.Context, &wallet.Config{
    94  		WalletDir:      cfg.walletDir,
    95  		WalletPassword: password,
    96  	})
    97  	assert.NoError(t, err)
    98  }
    99  
   100  func TestRecoverDerivedWallet_AlreadyExists(t *testing.T) {
   101  	cfg := setupRecoverCfg(t)
   102  	cfg.numAccounts = 4
   103  	cliCtx := createRecoverCliCtx(t, cfg)
   104  	require.NoError(t, RecoverWalletCli(cliCtx))
   105  
   106  	// Trying to recover an HD wallet into a directory that already exists should give an error
   107  	require.ErrorContains(t, "a wallet already exists at this location", RecoverWalletCli(cliCtx))
   108  }