From 5d7f177811e2f2f14b37671888398697251a7262 Mon Sep 17 00:00:00 2001
From: Evgenii Stratonikov <evgeniy@nspcc.ru>
Date: Mon, 31 Aug 2020 16:17:44 +0300
Subject: [PATCH] cli: add tests for wallet/account creation

---
 cli/wallet/wallet.go |  4 ++-
 cli/wallet_test.go   | 85 ++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 88 insertions(+), 1 deletion(-)
 create mode 100644 cli/wallet_test.go

diff --git a/cli/wallet/wallet.go b/cli/wallet/wallet.go
index c74e7a17b..e8f8e2389 100644
--- a/cli/wallet/wallet.go
+++ b/cli/wallet/wallet.go
@@ -474,7 +474,9 @@ func importWallet(ctx *cli.Context) error {
 		acc.Contract.Script = ctr
 	}
 
-	acc.Label = ctx.String("name")
+	if acc.Label == "" {
+		acc.Label = ctx.String("name")
+	}
 	if err := addAccountAndSave(wall, acc); err != nil {
 		return cli.NewExitError(err, 1)
 	}
diff --git a/cli/wallet_test.go b/cli/wallet_test.go
new file mode 100644
index 000000000..2a57a2fdf
--- /dev/null
+++ b/cli/wallet_test.go
@@ -0,0 +1,85 @@
+package main
+
+import (
+	"os"
+	"path"
+	"testing"
+
+	"github.com/nspcc-dev/neo-go/pkg/crypto/keys"
+	"github.com/nspcc-dev/neo-go/pkg/wallet"
+	"github.com/stretchr/testify/require"
+)
+
+func TestWalletInit(t *testing.T) {
+	tmpDir := os.TempDir()
+	e := newExecutor(t, false)
+	defer e.Close(t)
+
+	walletPath := path.Join(tmpDir, "wallet.json")
+	e.Run(t, "neo-go", "wallet", "init", "--wallet", walletPath)
+	defer os.Remove(walletPath)
+
+	t.Run("CreateAccount", func(t *testing.T) {
+		e.In.WriteString("testname\r")
+		e.In.WriteString("testpass\r")
+		e.In.WriteString("testpass\r")
+		e.Run(t, "neo-go", "wallet", "create", "--wallet", walletPath)
+
+		w, err := wallet.NewWalletFromFile(walletPath)
+		require.NoError(t, err)
+		defer w.Close()
+		require.Len(t, w.Accounts, 1)
+		require.Equal(t, w.Accounts[0].Label, "testname")
+		require.NoError(t, w.Accounts[0].Decrypt("testpass"))
+	})
+
+	t.Run("Import", func(t *testing.T) {
+		t.Run("WIF", func(t *testing.T) {
+			priv, err := keys.NewPrivateKey()
+			require.NoError(t, err)
+			e.In.WriteString("test_account\r")
+			e.In.WriteString("qwerty\r")
+			e.In.WriteString("qwerty\r")
+			e.Run(t, "neo-go", "wallet", "import", "--wallet", walletPath,
+				"--wif", priv.WIF())
+
+			w, err := wallet.NewWalletFromFile(walletPath)
+			require.NoError(t, err)
+			defer w.Close()
+			acc := w.GetAccount(priv.GetScriptHash())
+			require.NotNil(t, acc)
+			require.Equal(t, "test_account", acc.Label)
+			require.NoError(t, acc.Decrypt("qwerty"))
+
+			t.Run("AlreadyExists", func(t *testing.T) {
+				e.In.WriteString("test_account_2\r")
+				e.In.WriteString("qwerty2\r")
+				e.In.WriteString("qwerty2\r")
+				e.RunWithError(t, "neo-go", "wallet", "import",
+					"--wallet", walletPath, "--wif", priv.WIF())
+			})
+		})
+		t.Run("EncryptedWIF", func(t *testing.T) {
+			acc, err := wallet.NewAccount()
+			require.NoError(t, err)
+			require.NoError(t, acc.Encrypt("somepass"))
+
+			t.Run("InvalidPassword", func(t *testing.T) {
+				e.In.WriteString("password1\r")
+				e.RunWithError(t, "neo-go", "wallet", "import", "--wallet", walletPath,
+					"--wif", acc.EncryptedWIF)
+			})
+
+			e.In.WriteString("somepass\r")
+			e.Run(t, "neo-go", "wallet", "import", "--wallet", walletPath,
+				"--wif", acc.EncryptedWIF)
+
+			w, err := wallet.NewWalletFromFile(walletPath)
+			require.NoError(t, err)
+			defer w.Close()
+			actual := w.GetAccount(acc.PrivateKey().GetScriptHash())
+			require.NotNil(t, actual)
+			require.NoError(t, actual.Decrypt("somepass"))
+		})
+	})
+}