diff --git a/pkg/wallet/wallet.go b/pkg/wallet/wallet.go index 447df7bff..1b2fc731e 100644 --- a/pkg/wallet/wallet.go +++ b/pkg/wallet/wallet.go @@ -2,6 +2,7 @@ package wallet import ( "encoding/json" + "errors" "io" "os" @@ -102,6 +103,19 @@ func (w *Wallet) AddAccount(acc *Account) { w.Accounts = append(w.Accounts, acc) } +// RemoveAccount removes an Account with the specified addr +// from the wallet. +func (w *Wallet) RemoveAccount(addr string) error { + for i, acc := range w.Accounts { + if acc.Address == addr { + copy(w.Accounts[i:], w.Accounts[i+1:]) + w.Accounts = w.Accounts[:len(w.Accounts)-1] + return nil + } + } + return errors.New("account wasn't found") +} + // AddToken adds new token to a wallet. func (w *Wallet) AddToken(tok *Token) { w.Extra.Tokens = append(w.Extra.Tokens, tok) diff --git a/pkg/wallet/wallet_test.go b/pkg/wallet/wallet_test.go index 73e417ec3..128fbd4fe 100644 --- a/pkg/wallet/wallet_test.go +++ b/pkg/wallet/wallet_test.go @@ -48,7 +48,7 @@ func TestAddAccount(t *testing.T) { privateKey: nil, publicKey: nil, wif: "", - Address: "", + Address: "real", EncryptedWIF: "", Label: "", Contract: nil, @@ -57,6 +57,11 @@ func TestAddAccount(t *testing.T) { }) accounts := wallet.Accounts require.Len(t, accounts, 1) + + require.Error(t, wallet.RemoveAccount("abc")) + require.Len(t, wallet.Accounts, 1) + require.NoError(t, wallet.RemoveAccount("real")) + require.Len(t, wallet.Accounts, 0) } func TestPath(t *testing.T) {