From fdcc72dad0868af7a7482abd8287a2fa9c74ffea Mon Sep 17 00:00:00 2001
From: Roman Khimov <roman@nspcc.ru>
Date: Mon, 28 Sep 2020 17:24:21 +0300
Subject: [PATCH] cli/wallet: make 'addr' optional for NEP5 balances query

Iterate over all accounts by default.
---
 cli/executor_test.go |  7 +++-
 cli/nep5_test.go     | 61 ++++++++++++++++++++++++----
 cli/wallet/nep5.go   | 96 +++++++++++++++++++++++++-------------------
 3 files changed, 115 insertions(+), 49 deletions(-)

diff --git a/cli/executor_test.go b/cli/executor_test.go
index 2a753d02b..e4198d437 100644
--- a/cli/executor_test.go
+++ b/cli/executor_test.go
@@ -127,9 +127,14 @@ func (e *executor) GetTransaction(t *testing.T, h util.Uint256) (*transaction.Tr
 	return tx, height
 }
 
-func (e *executor) checkNextLine(t *testing.T, expected string) {
+func (e *executor) getNextLine(t *testing.T) string {
 	line, err := e.Out.ReadString('\n')
 	require.NoError(t, err)
+	return line
+}
+
+func (e *executor) checkNextLine(t *testing.T, expected string) {
+	line := e.getNextLine(t)
 	e.checkLine(t, line, expected)
 }
 
diff --git a/cli/nep5_test.go b/cli/nep5_test.go
index c99f3cda1..57d1185a8 100644
--- a/cli/nep5_test.go
+++ b/cli/nep5_test.go
@@ -6,6 +6,7 @@ import (
 	"os"
 	"path"
 	"strconv"
+	"strings"
 	"testing"
 
 	"github.com/nspcc-dev/neo-go/pkg/encoding/address"
@@ -18,15 +19,16 @@ import (
 func TestNEP5Balance(t *testing.T) {
 	e := newExecutor(t, true)
 	defer e.Close(t)
-	cmd := []string{
-		"neo-go", "wallet", "nep5", "balance",
-		"--rpc-endpoint", "http://" + e.RPC.Addr,
+	cmdbalance := []string{"neo-go", "wallet", "nep5", "balance"}
+	cmdbase := append(cmdbalance,
+		"--rpc-endpoint", "http://"+e.RPC.Addr,
 		"--wallet", validatorWallet,
-		"--addr", validatorAddr,
-	}
+	)
+	cmd := append(cmdbase, "--addr", validatorAddr)
 	t.Run("NEO", func(t *testing.T) {
 		b, index := e.Chain.GetGoverningTokenBalance(validatorHash)
 		checkResult := func(t *testing.T) {
+			e.checkNextLine(t, "^\\s*Account\\s+"+validatorAddr)
 			e.checkNextLine(t, "^\\s*NEO:\\s+NEO \\("+e.Chain.GoverningTokenHash().StringLE()+"\\)")
 			e.checkNextLine(t, "^\\s*Amount\\s*:\\s*"+b.String())
 			e.checkNextLine(t, "^\\s*Updated\\s*:\\s*"+strconv.FormatUint(uint64(index), 10))
@@ -43,12 +45,57 @@ func TestNEP5Balance(t *testing.T) {
 	})
 	t.Run("GAS", func(t *testing.T) {
 		e.Run(t, append(cmd, "--token", "gas")...)
+		e.checkNextLine(t, "^\\s*Account\\s+"+validatorAddr)
 		e.checkNextLine(t, "^\\s*GAS:\\s+GAS \\("+e.Chain.UtilityTokenHash().StringLE()+"\\)")
 		b := e.Chain.GetUtilityTokenBalance(validatorHash)
 		e.checkNextLine(t, "^\\s*Amount\\s*:\\s*"+util.Fixed8(b.Int64()).String())
 	})
-	t.Run("Invalid", func(t *testing.T) {
-		e.RunWithError(t, append(cmd, "--token", "kek")...)
+	t.Run("all accounts", func(t *testing.T) {
+		e.Run(t, cmdbase...)
+		addr1, err := address.StringToUint160("NbTiM6h8r99kpRtb428XcsUk1TzKed2gTc")
+		require.NoError(t, err)
+		e.checkNextLine(t, "^Account "+address.Uint160ToString(addr1))
+		e.checkNextLine(t, "^\\s*GAS:\\s+GAS \\("+e.Chain.UtilityTokenHash().StringLE()+"\\)")
+		balance := e.Chain.GetUtilityTokenBalance(addr1)
+		e.checkNextLine(t, "^\\s*Amount\\s*:\\s*"+util.Fixed8(balance.Int64()).String())
+		e.checkNextLine(t, "^\\s*Updated:")
+		e.checkNextLine(t, "^\\s*$")
+
+		addr2, err := address.StringToUint160("NUVPACMnKFhpuHjsRjhUvXz1XhqfGZYVtY")
+		require.NoError(t, err)
+		e.checkNextLine(t, "^Account "+address.Uint160ToString(addr2))
+		e.checkNextLine(t, "^\\s*$")
+
+		addr3, err := address.StringToUint160("NVNvVRW5Q5naSx2k2iZm7xRgtRNGuZppAK")
+		require.NoError(t, err)
+		e.checkNextLine(t, "^Account "+address.Uint160ToString(addr3))
+		// The order of assets is undefined.
+		for i := 0; i < 2; i++ {
+			line := e.getNextLine(t)
+			if strings.Contains(line, "GAS") {
+				e.checkLine(t, line, "^\\s*GAS:\\s+GAS \\("+e.Chain.UtilityTokenHash().StringLE()+"\\)")
+				balance = e.Chain.GetUtilityTokenBalance(addr3)
+				e.checkNextLine(t, "^\\s*Amount\\s*:\\s*"+util.Fixed8(balance.Int64()).String())
+				e.checkNextLine(t, "^\\s*Updated:")
+			} else {
+				balance, index := e.Chain.GetGoverningTokenBalance(validatorHash)
+				e.checkLine(t, line, "^\\s*NEO:\\s+NEO \\("+e.Chain.GoverningTokenHash().StringLE()+"\\)")
+				e.checkNextLine(t, "^\\s*Amount\\s*:\\s*"+balance.String())
+				e.checkNextLine(t, "^\\s*Updated\\s*:\\s*"+strconv.FormatUint(uint64(index), 10))
+			}
+		}
+		e.checkEOF(t)
+	})
+	t.Run("Bad token", func(t *testing.T) {
+		e.Run(t, append(cmd, "--token", "kek")...)
+		e.checkNextLine(t, "^\\s*Account\\s+"+validatorAddr)
+		e.checkEOF(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, "--addr", "xxx")...)
 	})
 	return
 }
diff --git a/cli/wallet/nep5.go b/cli/wallet/nep5.go
index 106857e6a..727be16e6 100644
--- a/cli/wallet/nep5.go
+++ b/cli/wallet/nep5.go
@@ -78,7 +78,7 @@ func newNEP5Commands() []cli.Command {
 		{
 			Name:      "balance",
 			Usage:     "get address balance",
-			UsageText: "balance --wallet <path> --rpc-endpoint <node> --timeout <time> --addr <addr> [--token <hash-or-name>]",
+			UsageText: "balance --wallet <path> --rpc-endpoint <node> [--timeout <time>] [--addr <addr>] [--token <hash-or-name>]",
 			Action:    getNEP5Balance,
 			Flags:     balanceFlags,
 		},
@@ -135,20 +135,30 @@ func newNEP5Commands() []cli.Command {
 }
 
 func getNEP5Balance(ctx *cli.Context) error {
+	var accounts []*wallet.Account
+
 	wall, err := openWallet(ctx.String("wallet"))
 	if err != nil {
-		return cli.NewExitError(err, 1)
+		return cli.NewExitError(fmt.Errorf("bad wallet: %w", err), 1)
 	}
 	defer wall.Close()
 
 	addr := ctx.String("addr")
-	addrHash, err := address.StringToUint160(addr)
-	if err != nil {
-		return cli.NewExitError(fmt.Errorf("invalid address: %w", err), 1)
-	}
-	acc := wall.GetAccount(addrHash)
-	if acc == nil {
-		return cli.NewExitError(fmt.Errorf("can't find account for the address: %s", addr), 1)
+	if addr != "" {
+		addrHash, err := address.StringToUint160(addr)
+		if err != nil {
+			return cli.NewExitError(fmt.Errorf("invalid address: %w", err), 1)
+		}
+		acc := wall.GetAccount(addrHash)
+		if acc == nil {
+			return cli.NewExitError(fmt.Errorf("can't find account for the address: %s", addr), 1)
+		}
+		accounts = append(accounts, acc)
+	} else {
+		if len(wall.Accounts) == 0 {
+			return cli.NewExitError(errors.New("no accounts in the wallet"), 1)
+		}
+		accounts = wall.Accounts
 	}
 
 	gctx, cancel := options.GetTimeoutContext(ctx)
@@ -159,44 +169,48 @@ func getNEP5Balance(ctx *cli.Context) error {
 		return err
 	}
 
-	var token *wallet.Token
 	name := ctx.String("token")
-	if name != "" {
-		token, err = getMatchingToken(ctx, wall, name)
+
+	for k, acc := range accounts {
+		addrHash, err := address.StringToUint160(acc.Address)
 		if err != nil {
-			token, err = getMatchingTokenRPC(ctx, c, addrHash, name)
+			return cli.NewExitError(fmt.Errorf("invalid account address: %w", err), 1)
+		}
+		balances, err := c.GetNEP5Balances(addrHash)
+		if err != nil {
+			return cli.NewExitError(err, 1)
+		}
+
+		if k != 0 {
+			fmt.Fprintln(ctx.App.Writer)
+		}
+		fmt.Fprintf(ctx.App.Writer, "Account %s\n", acc.Address)
+
+		for i := range balances.Balances {
+			var tokenName, tokenSymbol string
+
+			asset := balances.Balances[i].Asset
+			token, err := getMatchingToken(ctx, wall, asset.StringLE())
 			if err != nil {
-				return cli.NewExitError(err, 1)
+				token, err = c.NEP5TokenInfo(asset)
 			}
+			if err == nil {
+				if name != "" && !(token.Name == name || token.Symbol == name || token.Address() == name || token.Hash.StringLE() == name) {
+					continue
+				}
+				tokenName = token.Name
+				tokenSymbol = token.Symbol
+			} else {
+				if name != "" {
+					continue
+				}
+				tokenSymbol = "UNKNOWN"
+			}
+			fmt.Fprintf(ctx.App.Writer, "%s: %s (%s)\n", strings.ToUpper(tokenSymbol), tokenName, asset.StringLE())
+			fmt.Fprintf(ctx.App.Writer, "\tAmount : %s\n", balances.Balances[i].Amount)
+			fmt.Fprintf(ctx.App.Writer, "\tUpdated: %d\n", balances.Balances[i].LastUpdated)
 		}
 	}
-
-	balances, err := c.GetNEP5Balances(addrHash)
-	if err != nil {
-		return cli.NewExitError(err, 1)
-	}
-
-	for i := range balances.Balances {
-		var tokenName, tokenSymbol string
-
-		asset := balances.Balances[i].Asset
-		if name != "" && !token.Hash.Equals(asset) {
-			continue
-		}
-		token, err := getMatchingToken(ctx, wall, asset.StringLE())
-		if err != nil {
-			token, err = c.NEP5TokenInfo(asset)
-		}
-		if err == nil {
-			tokenName = token.Name
-			tokenSymbol = token.Symbol
-		} else {
-			tokenSymbol = "UNKNOWN"
-		}
-		fmt.Fprintf(ctx.App.Writer, "%s: %s (%s)\n", strings.ToUpper(tokenSymbol), tokenName, asset.StringLE())
-		fmt.Fprintf(ctx.App.Writer, "\tAmount : %s\n", balances.Balances[i].Amount)
-		fmt.Fprintf(ctx.App.Writer, "\tUpdated: %d\n", balances.Balances[i].LastUpdated)
-	}
 	return nil
 }