Merge branch 'master' into illia-malachyn/744-account-keys-endpoint

onflow · Sep 18, 2024 · bed1843 · bed1843
2 parents a1405b5 + 356f7b7
commit bed1843
Show file tree

Hide file tree

Showing 3 changed files with 111 additions and 0 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) GetAccountKeyAtLatestBlock(ctx context.Context, address flow.Address, keyIndex uint32) (*flow.AccountKey, error) {
 	return c.grpc.GetAccountKeyAtLatestBlock(ctx, address, keyIndex)
 }

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) GetAccountKeyAtLatestBlock(
 	ctx context.Context,
 	address flow.Address,

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)