wallet: support token removal

Also add tests for Add/RemoveToken
This commit is contained in:
Evgenii Stratonikov 2020-03-13 17:23:45 +03:00
parent cc5ec8e6d6
commit 9a41ffb9fb
2 changed files with 25 additions and 0 deletions

View file

@ -121,6 +121,18 @@ func (w *Wallet) AddToken(tok *Token) {
w.Extra.Tokens = append(w.Extra.Tokens, tok)
}
// RemoveToken removes token with the specified hash from the wallet.
func (w *Wallet) RemoveToken(h util.Uint160) error {
for i, tok := range w.Extra.Tokens {
if tok.Hash.Equals(h) {
copy(w.Extra.Tokens[i:], w.Extra.Tokens[i+1:])
w.Extra.Tokens = w.Extra.Tokens[:len(w.Extra.Tokens)-1]
return nil
}
}
return errors.New("token wasn't found")
}
// Path returns the location of the wallet on the filesystem.
func (w *Wallet) Path() string {
return w.path

View file

@ -7,6 +7,7 @@ import (
"testing"
"github.com/nspcc-dev/neo-go/pkg/encoding/address"
"github.com/nspcc-dev/neo-go/pkg/util"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)
@ -138,6 +139,18 @@ func removeWallet(t *testing.T, walletPath string) {
require.NoError(t, err)
}
func TestWallet_AddToken(t *testing.T) {
w := checkWalletConstructor(t)
tok := NewToken(util.Uint160{1, 2, 3}, "Rubl", "RUB", 2)
require.Equal(t, 0, len(w.Extra.Tokens))
w.AddToken(tok)
require.Equal(t, 1, len(w.Extra.Tokens))
require.Error(t, w.RemoveToken(util.Uint160{4, 5, 6}))
require.Equal(t, 1, len(w.Extra.Tokens))
require.NoError(t, w.RemoveToken(tok.Hash))
require.Equal(t, 0, len(w.Extra.Tokens))
}
func TestWallet_GetAccount(t *testing.T) {
wallet := checkWalletConstructor(t)
accounts := []*Account{