Merge pull request #742 from The-K-R-O-K/illia-malachyn/741-get-accou…

…nt-balance-endpoint Add GetAccountBalance endpoint to grpc server
onflow · Sep 17, 2024 · 356f7b7 · 356f7b7
2 parents 160a1bc + 6e6b594
commit 356f7b7
Show file tree

Hide file tree

Showing 8 changed files with 330 additions and 6 deletions.
diff --git a/access/grpc/client.go b/access/grpc/client.go
@@ -200,6 +200,14 @@ func (c *Client) GetAccountAtBlockHeight(ctx context.Context, address flow.Addre
 	return c.grpc.GetAccountAtBlockHeight(ctx, address, blockHeight)
 }
 
+func (c *Client) GetAccountBalanceAtLatestBlock(ctx context.Context, address flow.Address) (uint64, error) {
+	return c.grpc.GetAccountBalanceAtLatestBlock(ctx, address)
+}
+
+func (c *Client) GetAccountBalanceAtBlockHeight(ctx context.Context, address flow.Address, blockHeight uint64) (uint64, error) {
+	return c.grpc.GetAccountBalanceAtBlockHeight(ctx, address, blockHeight)
+}
+
 func (c *Client) ExecuteScriptAtLatestBlock(ctx context.Context, script []byte, arguments []cadence.Value) (cadence.Value, error) {
 	return c.grpc.ExecuteScriptAtLatestBlock(ctx, script, arguments)
 }

diff --git a/access/grpc/grpc.go b/access/grpc/grpc.go
@@ -516,6 +516,42 @@ func (c *BaseClient) GetAccountAtBlockHeight(
 	return &account, nil
 }
 
+func (c *BaseClient) GetAccountBalanceAtLatestBlock(
+	ctx context.Context,
+	address flow.Address,
+	opts ...grpc.CallOption,
+) (uint64, error) {
+	request := &access.GetAccountBalanceAtLatestBlockRequest{
+		Address: address.Bytes(),
+	}
+
+	response, err := c.rpcClient.GetAccountBalanceAtLatestBlock(ctx, request, opts...)
+	if err != nil {
+		return 0, newRPCError(err)
+	}
+
+	return response.GetBalance(), nil
+}
+
+func (c *BaseClient) GetAccountBalanceAtBlockHeight(
+	ctx context.Context,
+	address flow.Address,
+	blockHeight uint64,
+	opts ...grpc.CallOption,
+) (uint64, error) {
+	request := &access.GetAccountBalanceAtBlockHeightRequest{
+		Address:     address.Bytes(),
+		BlockHeight: blockHeight,
+	}
+
+	response, err := c.rpcClient.GetAccountBalanceAtBlockHeight(ctx, request, opts...)
+	if err != nil {
+		return 0, newRPCError(err)
+	}
+
+	return response.GetBalance(), nil
+}
+
 func (c *BaseClient) ExecuteScriptAtLatestBlock(
 	ctx context.Context,
 	script []byte,

diff --git a/access/grpc/grpc_test.go b/access/grpc/grpc_test.go
@@ -709,6 +709,73 @@ func TestClient_GetAccountAtBlockHeight(t *testing.T) {
 	}))
 }
 
+func TestClient_GetAccountBalanceAtLatestBlock(t *testing.T) {
+	accounts := test.AccountGenerator()
+	addresses := test.AddressGenerator()
+
+	t.Run("Success", clientTest(func(t *testing.T, ctx context.Context, rpc *mocks.MockRPCClient, c *BaseClient) {
+		account := accounts.New()
+
+		response := &access.AccountBalanceResponse{
+			Balance: account.Balance,
+		}
+
+		rpc.On("GetAccountBalanceAtLatestBlock", ctx, mock.Anything).Return(response, nil)
+
+		balance, err := c.GetAccountBalanceAtLatestBlock(ctx, account.Address)
+		require.NoError(t, err)
+
+		assert.Equal(t, account.Balance, balance)
+
+	}))
+
+	t.Run("Not found error", clientTest(func(t *testing.T, ctx context.Context, rpc *mocks.MockRPCClient, c *BaseClient) {
+		address := addresses.New()
+
+		rpc.On("GetAccountBalanceAtLatestBlock", ctx, mock.Anything).
+			Return(nil, errNotFound)
+
+		balance, err := c.GetAccountBalanceAtLatestBlock(ctx, address)
+		assert.Error(t, err)
+		assert.Equal(t, codes.NotFound, status.Code(err))
+		assert.Equal(t, balance, uint64(0))
+	}))
+}
+
+func TestClient_GetAccountBalanceAtBlockHeight(t *testing.T) {
+	accounts := test.AccountGenerator()
+	addresses := test.AddressGenerator()
+	blockHeight := uint64(42)
+
+	t.Run("Success", clientTest(func(t *testing.T, ctx context.Context, rpc *mocks.MockRPCClient, c *BaseClient) {
+		account := accounts.New()
+
+		response := &access.AccountBalanceResponse{
+			Balance: account.Balance,
+		}
+
+		rpc.On("GetAccountBalanceAtBlockHeight", ctx, mock.Anything).Return(response, nil)
+
+		balance, err := c.GetAccountBalanceAtBlockHeight(ctx, account.Address, blockHeight)
+		require.NoError(t, err)
+
+		assert.Equal(t, account.Balance, balance)
+
+	}))
+
+	t.Run("Not found error", clientTest(func(t *testing.T, ctx context.Context, rpc *mocks.MockRPCClient, c *BaseClient) {
+		address := addresses.New()
+
+		rpc.On("GetAccountBalanceAtBlockHeight", ctx, mock.Anything).
+			Return(nil, errNotFound)
+
+		balance, err := c.GetAccountBalanceAtBlockHeight(ctx, address, blockHeight)
+		assert.Error(t, err)
+		assert.Equal(t, codes.NotFound, status.Code(err))
+		assert.Equal(t, balance, uint64(0))
+	}))
+}
+
 func TestClient_ExecuteScriptAtLatestBlock(t *testing.T) {
 	t.Run("Success", clientTest(func(t *testing.T, ctx context.Context, rpc *mocks.MockRPCClient, c *BaseClient) {
 		expectedValue := cadence.NewInt(42)

diff --git a/access/grpc/mocks/RPCClient.go b/access/grpc/mocks/RPCClient.go
diff --git a/examples/go.mod b/examples/go.mod
@@ -7,7 +7,7 @@ toolchain go1.22.4
 replace github.com/onflow/flow-go-sdk => ../
 
 require (
-	github.com/onflow/cadence v1.0.0-preview.38
+	github.com/onflow/cadence v1.0.0-preview.52
 	github.com/onflow/flow-cli/flowkit v1.11.0
 	github.com/onflow/flow-go-sdk v0.41.17
 	github.com/spf13/afero v1.11.0
@@ -41,9 +41,9 @@ require (
 	github.com/logrusorgru/aurora/v4 v4.0.0 // indirect
 	github.com/mattn/go-colorable v0.1.13 // indirect
 	github.com/mattn/go-isatty v0.0.20 // indirect
-	github.com/onflow/atree v0.7.0-rc.2 // indirect
+	github.com/onflow/atree v0.8.0-rc.6 // indirect
 	github.com/onflow/crypto v0.25.1 // indirect
-	github.com/onflow/flow/protobuf/go/flow v0.4.3 // indirect
+	github.com/onflow/flow/protobuf/go/flow v0.4.7 // indirect
 	github.com/onflow/sdks v0.6.0-preview.1 // indirect
 	github.com/pkg/errors v0.9.1 // indirect
 	github.com/pmezard/go-difflib v1.0.0 // indirect