From ff2b0516469d824f363b7e9282c2e3c195fd9045 Mon Sep 17 00:00:00 2001
From: Anna Shaleva <shaleva.ann@nspcc.ru>
Date: Fri, 16 Apr 2021 11:58:04 +0300
Subject: [PATCH] cli: use AddressFlag for nep17-related commands

---
 cli/nep17_test.go   |  3 ---
 cli/wallet/nep17.go | 13 +++++--------
 2 files changed, 5 insertions(+), 11 deletions(-)

diff --git a/cli/nep17_test.go b/cli/nep17_test.go
index 6930cfe85..20e5bb984 100644
--- a/cli/nep17_test.go
+++ b/cli/nep17_test.go
@@ -98,9 +98,6 @@ func TestNEP17Balance(t *testing.T) {
 	t.Run("Bad wallet", func(t *testing.T) {
 		e.RunWithError(t, append(cmdbalance, "--wallet", "/dev/null")...)
 	})
-	t.Run("Bad address", func(t *testing.T) {
-		e.RunWithError(t, append(cmdbalance, "--rpc-endpoint", "http://"+e.RPC.Addr, "--wallet", validatorWallet, "--address", "xxx")...)
-	})
 	return
 }
 
diff --git a/cli/wallet/nep17.go b/cli/wallet/nep17.go
index 76096bb94..09ba5933c 100644
--- a/cli/wallet/nep17.go
+++ b/cli/wallet/nep17.go
@@ -32,7 +32,7 @@ func newNEP17Commands() []cli.Command {
 	balanceFlags := []cli.Flag{
 		walletPathFlag,
 		tokenFlag,
-		cli.StringFlag{
+		flags.AddressFlag{
 			Name:  "address, a",
 			Usage: "Address to use",
 		},
@@ -135,15 +135,12 @@ func getNEP17Balance(ctx *cli.Context) error {
 	}
 	defer wall.Close()
 
-	addr := ctx.String("address")
-	if addr != "" {
-		addrHash, err := address.StringToUint160(addr)
-		if err != nil {
-			return cli.NewExitError(fmt.Errorf("invalid address: %w", err), 1)
-		}
+	addrFlag := ctx.Generic("address").(*flags.Address)
+	if addrFlag.IsSet {
+		addrHash := addrFlag.Uint160()
 		acc := wall.GetAccount(addrHash)
 		if acc == nil {
-			return cli.NewExitError(fmt.Errorf("can't find account for the address: %s", addr), 1)
+			return cli.NewExitError(fmt.Errorf("can't find account for the address: %s", address.Uint160ToString(addrHash)), 1)
 		}
 		accounts = append(accounts, acc)
 	} else {