From 6af2f3fdecde11fb89728df6b6d0f7983c5284a1 Mon Sep 17 00:00:00 2001
From: Anna Shaleva <shaleva.ann@nspcc.ru>
Date: Mon, 19 Apr 2021 16:45:40 +0300
Subject: [PATCH] cli: use AddressFlag for 'wallet remove' command

---
 cli/wallet/wallet.go | 17 ++++++++++-------
 cli/wallet_test.go   |  2 +-
 2 files changed, 11 insertions(+), 8 deletions(-)

diff --git a/cli/wallet/wallet.go b/cli/wallet/wallet.go
index c263dfff7..ad7ee8489 100644
--- a/cli/wallet/wallet.go
+++ b/cli/wallet/wallet.go
@@ -206,11 +206,15 @@ func NewCommands() []cli.Command {
 			{
 				Name:      "remove",
 				Usage:     "remove an account from the wallet",
-				UsageText: "remove --wallet <path> [--force] <addr>",
+				UsageText: "remove --wallet <path> [--force] --address <addr>",
 				Action:    removeAccount,
 				Flags: []cli.Flag{
 					walletPathFlag,
 					forceFlag,
+					flags.AddressFlag{
+						Name:  "address, a",
+						Usage: "Account address or hash in LE form to be removed",
+					},
 				},
 			},
 			{
@@ -502,18 +506,17 @@ func removeAccount(ctx *cli.Context) error {
 	}
 	defer wall.Close()
 
-	addrArg := ctx.Args().First()
-	addr, err := address.StringToUint160(addrArg)
-	if err != nil {
-		return cli.NewExitError("valid address must be provided", 1)
+	addr := ctx.Generic("address").(*flags.Address)
+	if !addr.IsSet {
+		cli.NewExitError("valid account address must be provided", 1)
 	}
-	acc := wall.GetAccount(addr)
+	acc := wall.GetAccount(addr.Uint160())
 	if acc == nil {
 		return cli.NewExitError("account wasn't found", 1)
 	}
 
 	if !ctx.Bool("force") {
-		fmt.Fprintf(ctx.App.Writer, "Account %s will be removed. This action is irreversible.\n", addrArg)
+		fmt.Fprintf(ctx.App.Writer, "Account %s will be removed. This action is irreversible.\n", addr.Uint160())
 		if ok := askForConsent(ctx.App.Writer); !ok {
 			return nil
 		}
diff --git a/cli/wallet_test.go b/cli/wallet_test.go
index e780e32f6..d91d49448 100644
--- a/cli/wallet_test.go
+++ b/cli/wallet_test.go
@@ -71,7 +71,7 @@ func TestWalletInit(t *testing.T) {
 			addr := w.Accounts[0].Address
 			e.In.WriteString("y\r")
 			e.Run(t, "neo-go", "wallet", "remove",
-				"--wallet", walletPath, addr)
+				"--wallet", walletPath, "--address", addr)
 			w, err := wallet.NewWalletFromFile(walletPath)
 			require.NoError(t, err)
 			require.Nil(t, w.GetAccount(sh))