From ddf1ac0f289bf10aed1723db54897786dd35803c Mon Sep 17 00:00:00 2001
From: Alex Vanin <alexey@nspcc.ru>
Date: Thu, 25 Feb 2021 19:50:10 +0300
Subject: [PATCH] [#404] morph/client: Add awaiting function

Awaiting function locks execution for N blocks.
Useful to wait for notary deposit.

Signed-off-by: Alex Vanin <alexey@nspcc.ru>
---
 pkg/morph/client/client.go      | 47 +++++++++++++++++++++++++++++++++
 pkg/morph/client/constructor.go | 23 ++++++++++------
 2 files changed, 62 insertions(+), 8 deletions(-)

diff --git a/pkg/morph/client/client.go b/pkg/morph/client/client.go
index 3fa47add3..27398bf02 100644
--- a/pkg/morph/client/client.go
+++ b/pkg/morph/client/client.go
@@ -1,6 +1,9 @@
 package client
 
 import (
+	"context"
+	"time"
+
 	"github.com/nspcc-dev/neo-go/pkg/core/transaction"
 	"github.com/nspcc-dev/neo-go/pkg/encoding/fixedn"
 	"github.com/nspcc-dev/neo-go/pkg/rpc/client"
@@ -29,6 +32,8 @@ type Client struct {
 
 	gas util.Uint160 // native gas script-hash
 
+	waitInterval time.Duration
+
 	notary *notary
 }
 
@@ -136,6 +141,48 @@ func (c *Client) TransferGas(receiver util.Uint160, amount fixedn.Fixed8) error
 	return nil
 }
 
+// Wait function blocks routing execution until there
+// are `n` new blocks in the chain.
+func (c *Client) Wait(ctx context.Context, n uint32) {
+	var (
+		err               error
+		height, newHeight uint32
+	)
+
+	height, err = c.client.GetBlockCount()
+	if err != nil {
+		c.logger.Error("can't get blockchain height",
+			zap.String("error", err.Error()))
+		return
+	}
+
+	for {
+		select {
+		case <-ctx.Done():
+			return
+		default:
+		}
+
+		newHeight, err = c.client.GetBlockCount()
+		if err != nil {
+			c.logger.Error("can't get blockchain height",
+				zap.String("error", err.Error()))
+			return
+		}
+
+		if newHeight >= height+n {
+			return
+		}
+
+		time.Sleep(c.waitInterval)
+	}
+}
+
+// GasBalance returns GAS amount in the client's wallet.
+func (c *Client) GasBalance() (int64, error) {
+	return c.client.NEP17BalanceOf(c.gas, c.acc.PrivateKey().GetScriptHash())
+}
+
 func toStackParameter(value interface{}) (sc.Parameter, error) {
 	var result = sc.Parameter{
 		Value: value,
diff --git a/pkg/morph/client/constructor.go b/pkg/morph/client/constructor.go
index b287feeee..7aaee6fd2 100644
--- a/pkg/morph/client/constructor.go
+++ b/pkg/morph/client/constructor.go
@@ -27,15 +27,21 @@ type cfg struct {
 	logger *logger.Logger // logging component
 
 	gas util.Uint160 // native gas script-hash
+
+	waitInterval time.Duration
 }
 
-const defaultDialTimeout = 5 * time.Second
+const (
+	defaultDialTimeout  = 5 * time.Second
+	defaultWaitInterval = 500 * time.Millisecond
+)
 
 func defaultConfig() *cfg {
 	return &cfg{
-		ctx:         context.Background(),
-		dialTimeout: defaultDialTimeout,
-		logger:      zap.L(),
+		ctx:          context.Background(),
+		dialTimeout:  defaultDialTimeout,
+		logger:       zap.L(),
+		waitInterval: defaultWaitInterval,
 	}
 }
 
@@ -95,10 +101,11 @@ func New(key *ecdsa.PrivateKey, endpoint string, opts ...Option) (*Client, error
 	}
 
 	return &Client{
-		logger: cfg.logger,
-		client: cli,
-		acc:    account,
-		gas:    gas,
+		logger:       cfg.logger,
+		client:       cli,
+		acc:          account,
+		gas:          gas,
+		waitInterval: cfg.waitInterval,
 	}, nil
 }