add the ability to query a contract balance

account_balance_query.go

@@ -2,11 +2,17 @@ package hedera
 
 import "github.com/hashgraph/hedera-sdk-go/proto"
 
+// AccountBalanceQuery gets the balance of a CryptoCurrency account. This returns only the balance, so it is a smaller
+// and faster reply than AccountInfoQuery, which returns the balance plus additional information.
 type AccountBalanceQuery struct {
 	QueryBuilder
 	pb *proto.CryptoGetAccountBalanceQuery
 }
 
+// NewAccountBalanceQuery creates an AccountBalanceQuery builder which can be used to construct and execute
+// an AccountBalanceQuery.
+// It is recommended that you use this for creating new instances of an AccountBalanceQuery
+// instead of manually creating an instance of the struct.
 func NewAccountBalanceQuery() *AccountBalanceQuery {
 	pb := &proto.CryptoGetAccountBalanceQuery{Header: &proto.QueryHeader{}}
 
@@ -16,11 +22,31 @@ func NewAccountBalanceQuery() *AccountBalanceQuery {
 	return &AccountBalanceQuery{inner, pb}
 }
 
+// SetAccountID sets the AccountID for which you wish to query the balance.
+//
+// Note: you can only query an Account or Contract but not both -- if a Contract ID or Account ID has already been set,
+// it will be overwritten by this method.
 func (builder *AccountBalanceQuery) SetAccountID(id AccountID) *AccountBalanceQuery {
-	builder.pb.AccountID = id.toProto()
+	builder.pb.BalanceSource = &proto.CryptoGetAccountBalanceQuery_AccountID{
+		AccountID: id.toProto(),
+	}
+
+	return builder
+}
+
+// SetContractID sets the ContractID for which you wish to query the balance.
+//
+// Note: you can only query an Account or Contract but not both -- if a Contract ID or Account ID has already been set,
+// it will be overwritten by this method.
+func (builder *AccountBalanceQuery) SetContractID(id ContractID) *AccountBalanceQuery {
+	builder.pb.BalanceSource = &proto.CryptoGetAccountBalanceQuery_ContractID{
+		ContractID: id.toProto(),
+	}
+
 	return builder
 }
 
+// Execute executes the AccountBalanceQuery using the provided client
 func (builder *AccountBalanceQuery) Execute(client *Client) (Hbar, error) {
 	resp, err := builder.execute(client)
 	if err != nil {

account_balance_query_test.go

@@ -16,3 +16,14 @@ func TestNewAccountBalanceQuery(t *testing.T) {
 
 	cupaloy.SnapshotT(t, query.pb.String())
 }
+
+func TestNewAccountBalanceQuery_ForContract(t *testing.T) {
+	mockTransaction, err := newMockTransaction()
+	assert.NoError(t, err)
+
+	query := NewAccountBalanceQuery().
+		SetContractID(ContractID{Contract: 3}).
+		SetQueryPaymentTransaction(mockTransaction)
+
+	cupaloy.SnapshotT(t, query.pb.String())
+}

proto/CryptoGetAccountBalance.proto

@@ -14,7 +14,10 @@ option go_package = "github.com/hashgraph/hedera-sdk-go/proto";
 /* Get the balance of a cryptocurrency account. This returns only the balance, so it is a smaller and faster reply than CryptoGetInfo, which returns the balance plus additional information. */
 message CryptoGetAccountBalanceQuery {
     QueryHeader header = 1; // Standard info sent from client to node, including the signed payment, and what kind of response is requested (cost, state proof, both, or neither).
-    AccountID accountID = 2; // The account ID for which information is requested
+    oneof balanceSource {
+        AccountID accountID = 2; // The account ID for which information is requested
+        ContractID contractID = 3; // The account ID for which information is requested
+    }
 }
 
 /* Response when the client sends the node CryptoGetAccountBalanceQuery */
@@ -23,5 +26,3 @@ message CryptoGetAccountBalanceResponse {
     AccountID accountID = 2; // The account ID that is being described (this is useful with state proofs, for proving to a third party)
     uint64 balance = 3; // The current balance, in tinybars
 }
-
-