neo-go/pkg/wallet/wallet_test.go

package wallet

import (
	"encoding/json"
	"io/ioutil"
	"os"
	"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"
)

const (
	walletTemplate = "testWallet"
)

func TestNewWallet(t *testing.T) {
	wallet := checkWalletConstructor(t)
	require.NotNil(t, wallet)
}

func TestNewWalletFromFile_Negative_EmptyFile(t *testing.T) {
	_ = checkWalletConstructor(t)
	walletFromFile, err2 := NewWalletFromFile(walletTemplate)
	require.Errorf(t, err2, "EOF")
	require.Nil(t, walletFromFile)
}

func TestNewWalletFromFile_Negative_NoFile(t *testing.T) {
	_, err := NewWalletFromFile(walletTemplate)
	require.Errorf(t, err, "open testWallet: no such file or directory")
}

func TestCreateAccount(t *testing.T) {
	wallet := checkWalletConstructor(t)

	errAcc := wallet.CreateAccount("testName", "testPass")
	require.NoError(t, errAcc)
	accounts := wallet.Accounts
	require.Len(t, accounts, 1)
}

func TestAddAccount(t *testing.T) {
	wallet := checkWalletConstructor(t)

	wallet.AddAccount(&Account{
		privateKey:   nil,
		publicKey:    nil,
		wif:          "",
		Address:      "real",
		EncryptedWIF: "",
		Label:        "",
		Contract:     nil,
		Locked:       false,
		Default:      false,
	})
	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) {
	wallet := checkWalletConstructor(t)

	path := wallet.Path()
	require.NotEmpty(t, path)
}

func TestSave(t *testing.T) {
	file, err := ioutil.TempFile("", walletTemplate)
	require.NoError(t, err)
	wallet, err := NewWallet(file.Name())
	require.NoError(t, err)
	wallet.AddAccount(&Account{
		privateKey:   nil,
		publicKey:    nil,
		wif:          "",
		Address:      "",
		EncryptedWIF: "",
		Label:        "",
		Contract:     nil,
		Locked:       false,
		Default:      false,
	})

	t.Cleanup(func() {
		removeWallet(t, file.Name())
	})
	errForSave := wallet.Save()
	require.NoError(t, errForSave)

	openedWallet, err := NewWalletFromFile(wallet.path)
	require.NoError(t, err)
	require.Equal(t, wallet.Accounts, openedWallet.Accounts)

	t.Run("change and rewrite", func(t *testing.T) {
		err := openedWallet.CreateAccount("test", "pass")
		require.NoError(t, err)

		w2, err := NewWalletFromFile(openedWallet.path)
		require.NoError(t, err)
		require.Equal(t, 2, len(w2.Accounts))
		require.NoError(t, w2.Accounts[1].Decrypt("pass"))
		require.Equal(t, openedWallet.Accounts, w2.Accounts)
	})
}

func TestJSONMarshallUnmarshal(t *testing.T) {
	wallet := checkWalletConstructor(t)

	bytes, err := wallet.JSON()
	require.NoError(t, err)
	require.NotNil(t, bytes)

	unmarshalledWallet := &Wallet{}
	errUnmarshal := json.Unmarshal(bytes, unmarshalledWallet)

	require.NoError(t, errUnmarshal)
	require.Equal(t, wallet.Version, unmarshalledWallet.Version)
	require.Equal(t, wallet.Accounts, unmarshalledWallet.Accounts)
	require.Equal(t, wallet.Scrypt, unmarshalledWallet.Scrypt)
}

func checkWalletConstructor(t *testing.T) *Wallet {
	file, err := ioutil.TempFile("", walletTemplate)
	require.NoError(t, err)
	wallet, err := NewWallet(file.Name())
	t.Cleanup(func() {
		removeWallet(t, file.Name())
	})
	require.NoError(t, err)
	return wallet
}

func removeWallet(t *testing.T, walletPath string) {
	err := os.RemoveAll(walletPath)
	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{
		{
			Contract: &Contract{
				Script: []byte{0, 1, 2, 3},
			},
		},
		{
			Contract: &Contract{
				Script: []byte{3, 2, 1, 0},
			},
		},
	}

	for _, acc := range accounts {
		acc.Address = address.Uint160ToString(acc.Contract.ScriptHash())
		wallet.AddAccount(acc)
	}

	for i, acc := range accounts {
		h := acc.Contract.ScriptHash()
		assert.Equal(t, acc, wallet.GetAccount(h), "can't get %d account", i)
	}
}

func TestWalletGetChangeAddress(t *testing.T) {
	w1, err := NewWalletFromFile("testdata/wallet1.json")
	require.NoError(t, err)
	sh := w1.GetChangeAddress()
	// No default address, the first one is used.
	expected, err := address.StringToUint160("NTh9TnZTstvAePEYWDGLLxidBikJE24uTo")
	require.NoError(t, err)
	require.Equal(t, expected, sh)
	w2, err := NewWalletFromFile("testdata/wallet2.json")
	require.NoError(t, err)
	sh = w2.GetChangeAddress()
	// Default address.
	expected, err = address.StringToUint160("NUREbqw2kfbPgDeEz8Dac2QxntGGqqFMm7")
	require.NoError(t, err)
	require.Equal(t, expected, sh)
}