From 878626866ba7990757339c4d8009b4b41088fb02 Mon Sep 17 00:00:00 2001
From: Lukasz Zimnoch <lukasz.zimnoch@keep.network>
Date: Tue, 17 Nov 2020 17:37:45 +0100
Subject: [PATCH] Expose `BalanceAt` method

---
 pkg/chain/ethereum/ethutil/ethutil.go         |  6 ++++++
 pkg/chain/ethereum/ethutil/rate_limiter.go    | 14 ++++++++++++++
 .../ethereum/ethutil/rate_limiter_test.go     | 19 +++++++++++++++++++
 3 files changed, 39 insertions(+)

diff --git a/pkg/chain/ethereum/ethutil/ethutil.go b/pkg/chain/ethereum/ethutil/ethutil.go
index 9fb70f5..78f0975 100644
--- a/pkg/chain/ethereum/ethutil/ethutil.go
+++ b/pkg/chain/ethereum/ethutil/ethutil.go
@@ -28,6 +28,12 @@ type EthereumClient interface {
 	bind.ContractBackend
 	ethereum.ChainReader
 	ethereum.TransactionReader
+
+	BalanceAt(
+		ctx context.Context,
+		account common.Address,
+		blockNumber *big.Int,
+	) (*big.Int, error)
 }
 
 // AddressFromHex converts the passed string to a common.Address and returns it,
diff --git a/pkg/chain/ethereum/ethutil/rate_limiter.go b/pkg/chain/ethereum/ethutil/rate_limiter.go
index 585d304..2e35ff8 100644
--- a/pkg/chain/ethereum/ethutil/rate_limiter.go
+++ b/pkg/chain/ethereum/ethutil/rate_limiter.go
@@ -335,3 +335,17 @@ func (rl *rateLimiter) TransactionReceipt(
 
 	return rl.EthereumClient.TransactionReceipt(ctx, txHash)
 }
+
+func (rl *rateLimiter) BalanceAt(
+	ctx context.Context,
+	account common.Address,
+	blockNumber *big.Int,
+) (*big.Int, error) {
+	err := rl.acquirePermit()
+	if err != nil {
+		return nil, fmt.Errorf("cannot acquire rate limiter permit: [%v]", err)
+	}
+	defer rl.releasePermit()
+
+	return rl.EthereumClient.BalanceAt(ctx, account, blockNumber)
+}
diff --git a/pkg/chain/ethereum/ethutil/rate_limiter_test.go b/pkg/chain/ethereum/ethutil/rate_limiter_test.go
index b16788d..3611abb 100644
--- a/pkg/chain/ethereum/ethutil/rate_limiter_test.go
+++ b/pkg/chain/ethereum/ethutil/rate_limiter_test.go
@@ -470,6 +470,15 @@ func (mec *mockEthereumClient) TransactionReceipt(
 	return nil, nil
 }
 
+func (mec *mockEthereumClient) BalanceAt(
+	ctx context.Context,
+	account common.Address,
+	blockNumber *big.Int,
+) (*big.Int, error) {
+	mec.mockRequest()
+	return nil, nil
+}
+
 func getTests(
 	client EthereumClient,
 ) map[string]struct{ function func() error } {
@@ -639,5 +648,15 @@ func getTests(
 				return err
 			},
 		},
+		"test BalanceAt": {
+			function: func() error {
+				_, err := client.BalanceAt(
+					context.Background(),
+					common.Address{},
+					big.NewInt(0),
+				)
+				return err
+			},
+		},
 	}
 }