From 9817da7c4e3c9cc2db6548299daf2c69f25556fe Mon Sep 17 00:00:00 2001
From: Evgenii Stratonikov <evgeniy@nspcc.ru>
Date: Tue, 1 Sep 2020 15:40:32 +0300
Subject: [PATCH] cli: add tests for `nep5 import/remove`

Also unify `remove` and `import` arguments.
---
 cli/nep5_test.go   | 59 ++++++++++++++++++++++++++++++++++++++++++++++
 cli/wallet/nep5.go |  8 ++-----
 2 files changed, 61 insertions(+), 6 deletions(-)

diff --git a/cli/nep5_test.go b/cli/nep5_test.go
index 0a9013fa8..4e25f46e2 100644
--- a/cli/nep5_test.go
+++ b/cli/nep5_test.go
@@ -1,7 +1,10 @@
 package main
 
 import (
+	"io"
 	"math/big"
+	"os"
+	"path"
 	"strconv"
 	"strings"
 	"testing"
@@ -134,3 +137,59 @@ func TestNEP5MultiTransfer(t *testing.T) {
 	b, _ = e.Chain.GetGoverningTokenBalance(privs[2].GetScriptHash())
 	require.Equal(t, big.NewInt(13), b)
 }
+
+func TestNEP5ImportToken(t *testing.T) {
+	e := newExecutor(t, true)
+	defer e.Close(t)
+
+	tmpDir := os.TempDir()
+	walletPath := path.Join(tmpDir, "walletForImport.json")
+	defer os.Remove(walletPath)
+
+	e.Run(t, "neo-go", "wallet", "init", "--wallet", walletPath)
+	e.Run(t, "neo-go", "wallet", "nep5", "import",
+		"--rpc-endpoint", "http://"+e.RPC.Addr,
+		"--wallet", walletPath,
+		"--token", client.GasContractHash.StringLE())
+	e.Run(t, "neo-go", "wallet", "nep5", "import",
+		"--rpc-endpoint", "http://"+e.RPC.Addr,
+		"--wallet", walletPath,
+		"--token", client.NeoContractHash.StringLE())
+
+	t.Run("Info", func(t *testing.T) {
+		checkGASInfo := func(t *testing.T) {
+			e.checkNextLine(t, "^Name:\\s*GAS")
+			e.checkNextLine(t, "^Symbol:\\s*gas")
+			e.checkNextLine(t, "^Hash:\\s*"+client.GasContractHash.StringLE())
+			e.checkNextLine(t, "^Decimals:\\s*8")
+			e.checkNextLine(t, "^Address:\\s*"+address.Uint160ToString(client.GasContractHash))
+		}
+		t.Run("WithToken", func(t *testing.T) {
+			e.Run(t, "neo-go", "wallet", "nep5", "info",
+				"--wallet", walletPath, "--token", client.GasContractHash.StringLE())
+			checkGASInfo(t)
+		})
+		t.Run("NoToken", func(t *testing.T) {
+			e.Run(t, "neo-go", "wallet", "nep5", "info",
+				"--wallet", walletPath)
+			checkGASInfo(t)
+			_, err := e.Out.ReadString('\n')
+			require.NoError(t, err)
+			e.checkNextLine(t, "^Name:\\s*NEO")
+			e.checkNextLine(t, "^Symbol:\\s*neo")
+			e.checkNextLine(t, "^Hash:\\s*"+client.NeoContractHash.StringLE())
+			e.checkNextLine(t, "^Decimals:\\s*0")
+			e.checkNextLine(t, "^Address:\\s*"+address.Uint160ToString(client.NeoContractHash))
+		})
+		t.Run("Remove", func(t *testing.T) {
+			e.In.WriteString("y\r")
+			e.Run(t, "neo-go", "wallet", "nep5", "remove",
+				"--wallet", walletPath, "--token", client.NeoContractHash.StringLE())
+			e.Run(t, "neo-go", "wallet", "nep5", "info",
+				"--wallet", walletPath)
+			checkGASInfo(t)
+			_, err := e.Out.ReadString('\n')
+			require.Equal(t, err, io.EOF)
+		})
+	})
+}
diff --git a/cli/wallet/nep5.go b/cli/wallet/nep5.go
index 0fa1d6350..5ca1e293f 100644
--- a/cli/wallet/nep5.go
+++ b/cli/wallet/nep5.go
@@ -105,7 +105,7 @@ func newNEP5Commands() []cli.Command {
 		{
 			Name:      "remove",
 			Usage:     "remove NEP5 token from the wallet",
-			UsageText: "remove --wallet <path> <hash-or-name>",
+			UsageText: "remove --wallet <path> --token <hash-or-name>",
 			Action:    removeNEP5Token,
 			Flags: []cli.Flag{
 				walletPathFlag,
@@ -314,11 +314,7 @@ func removeNEP5Token(ctx *cli.Context) error {
 	}
 	defer wall.Close()
 
-	name := ctx.Args().First()
-	if name == "" {
-		return cli.NewExitError("token must be specified", 1)
-	}
-	token, err := getMatchingToken(ctx, wall, name)
+	token, err := getMatchingToken(ctx, wall, ctx.String("token"))
 	if err != nil {
 		return cli.NewExitError(err, 1)
 	}