feat: add contract_id support for CryptoGetAccountBalanceQuery (hiero-ledger#1389)

Signed-off-by: Antonio Ceppellini <antonio.ceppellini@gmail.com>
Signed-off-by: AntonioCeppellini <128388022+AntonioCeppellini@users.noreply.github.com>

Author: AntonioCeppellini

CHANGELOG.md

@@ -90,6 +90,7 @@ This changelog is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.
 - Added workflow to prevent assigning intermediate issues to contributors without prior Good First Issue completion (#1143).
 - Added `Client.from_env()` and network-specific factory methods (e.g., `Client.for_testnet()`) to simplify client initialization and reduce boilerplate. [[#1251](https://github.com/hiero-ledger/hiero-sdk-python/issues/1251)]
 - Improved unit test coverage for `TransactionId` class, covering parsing logic, hashing, and scheduled transactions.
+- Add contract_id support for CryptoGetAccountBalanceQuery([#1293](https://github.com/hiero-ledger/hiero-sdk-python/issues/1293))
 - Chained Good First Issue assignment with mentor assignment to bypass GitHub's anti-recursion protection - mentor assignment now occurs immediately after successful user assignment in the same workflow execution. (#1369)
 - Add GitHub Actions script and workflow for automatic spam list updates.
 - Added technical docstrings and hardening (set -euo pipefail) to the pr-check-test-files.sh script (#1336)

examples/contract/contract_balance_query.py

@@ -0,0 +1,101 @@
+"""
+Contract Balance Query Example
+
+This script demonstrates how to:
+1. Set up a client connection to the Hedera network
+2. Create a file containing contract bytecode
+3. Create a contract
+4. Query the contract balance using CryptoGetAccountBalanceQuery.set_contract_id()
+
+Run with:
+  uv run -m examples.contract.contract_balance_query
+  python -m examples.contract.contract_balance_query
+"""
+
+import os
+import sys
+from dotenv import load_dotenv
+
+from hiero_sdk_python import (
+    Client,
+    ContractCreateTransaction,
+    CryptoGetAccountBalanceQuery,
+    Hbar,
+)
+
+from hiero_sdk_python.contract.contract_id import ContractId
+
+from .contracts import SIMPLE_CONTRACT_BYTECODE
+from hiero_sdk_python.response_code import ResponseCode
+
+load_dotenv()
+
+
+def setup_client() -> Client:
+    print("Initializing client from environment variables...")
+    try:
+        client = Client.from_env()
+        print(f"✅ Success! Connected as operator: {client.operator_account_id}")
+    except Exception as e:
+        print(f"❌ Failed: {e}")
+        sys.exit(1)
+    
+    return client
+
+
+def create_contract(client: Client, initial_balance_tinybars: int) -> ContractId:
+    """Create a contract using the bytecode file and return its ContractId."""
+    bytecode = bytes.fromhex(SIMPLE_CONTRACT_BYTECODE)
+
+    receipt = (
+        ContractCreateTransaction()
+        .set_bytecode(bytecode)
+        .set_gas(2_000_000)
+        .set_initial_balance(initial_balance_tinybars)
+        .set_contract_memo("Contract for balance query example")
+        .execute(client)
+    )
+
+    status_code = ResponseCode(receipt.status)
+    status_name = status_code.name
+
+    if status_name == ResponseCode.SUCCESS.name:
+        print("✅ Transaction succeeded!")
+    elif status_code.is_unknown:
+        print(f"❓ Unknown transaction status: {status_name}")
+        sys.exit(1)
+    else:
+        print("❌ Transaction failed!")
+        sys.exit(1)
+
+    return receipt.contract_id
+
+
+def get_contract_balance(client: Client, contract_id: ContractId):
+    """Query contract balance using CryptoGetAccountBalanceQuery.set_contract_id()."""
+    print(f"Querying balance for contract {contract_id} ...")
+    balance = CryptoGetAccountBalanceQuery().set_contract_id(contract_id).execute(client)
+
+    print("✅ Balance retrieved successfully!")
+    print(f"  Contract: {contract_id}")
+    print(f"  Hbars: {balance.hbars}")
+    return balance
+
+
+def main():
+    try:
+        client = setup_client()
+
+        initial_balance_tinybars = Hbar(1)
+        contract_id = create_contract(client, initial_balance_tinybars.to_tinybars())
+
+        print(f"✅ Contract created with ID: {contract_id}")
+        get_contract_balance(client, contract_id)
+
+    except Exception as e:
+        print(f"❌Error: {e}", file=sys.stderr)
+        sys.exit(1)
+
+
+if __name__ == "__main__":
+    main()

src/hiero_sdk_python/query/account_balance_query.py

@@ -6,6 +6,8 @@
 from hiero_sdk_python.account.account_balance import AccountBalance
 from hiero_sdk_python.executable import _Method
 from hiero_sdk_python.channels import _Channel
+from hiero_sdk_python.contract.contract_id import ContractId
+
 
 class CryptoGetAccountBalanceQuery(Query):
     """
@@ -15,29 +17,61 @@ class CryptoGetAccountBalanceQuery(Query):
     including hbars and tokens.
     """
 
-    def __init__(self, account_id: Optional[AccountId] = None) -> None:
+    def __init__(
+        self,
+        account_id: Optional[AccountId] = None,
+        contract_id: Optional[ContractId] = None,
+    ) -> None:
         """
         Initializes a new instance of the CryptoGetAccountBalanceQuery class.
 
         Args:
             account_id (AccountId, optional): The ID of the account to retrieve the balance for.
+            contract_id (ContractId, optional): The ID of the contract to retrieve the balance for.
         """
         super().__init__()
-        self.account_id: Optional[AccountId] = account_id
+        self.account_id: Optional[AccountId] = None
+        self.contract_id: Optional[ContractId] = None
+
+        if account_id is not None:
+            self.set_account_id(account_id)
+        if contract_id is not None:
+            self.set_contract_id(contract_id)
 
     def set_account_id(self, account_id: AccountId) -> "CryptoGetAccountBalanceQuery":
         """
         Sets the account ID for which to retrieve the balance.
+        Resets to None the contract ID.
 
         Args:
             account_id (AccountId): The ID of the account.
 
         Returns:
             CryptoGetAccountBalanceQuery: The current instance for method chaining.
         """
+        if not isinstance(account_id, AccountId):
+            raise TypeError("account_id must be an AccountId.")
+        self.contract_id = None
         self.account_id = account_id
         return self
 
+    def set_contract_id(self, contract_id: ContractId) -> "CryptoGetAccountBalanceQuery":
+        """
+        Sets the contract ID for which to retrieve the balance.
+        Resets to None the account ID.
+
+        Args:
+            contract_id (ContractId): The ID of the contract.
+
+        Returns:
+            CryptoGetAccountBalanceQuery: The current instance for method chaining.
+        """
+        if not isinstance(contract_id, ContractId):
+            raise TypeError("contract_id must be a ContractId.")
+        self.account_id = None
+        self.contract_id = contract_id
+        return self
+
     def _make_request(self) -> query_pb2.Query:
         """
         Constructs the protobuf request for the account balance query.
@@ -46,22 +80,34 @@ def _make_request(self) -> query_pb2.Query:
             query_pb2.Query: The protobuf Query object containing the account balance query.
 
         Raises:
-            ValueError: If the account ID is not set.
+            ValueError: If both the account ID and contract ID are not set.
+            ValueError: If both the account ID and contract ID are set.
             AttributeError: If the Query protobuf structure is invalid.
             Exception: If any other error occurs during request construction.
         """
         try:
-            if not self.account_id:
-                raise ValueError("Account ID must be set before making the request.")
+            if not self.account_id and not self.contract_id:
+                raise ValueError("Either account_id or contract_id must be set before making the request.")
+
+            if self.account_id and self.contract_id:
+                raise ValueError("Specify either account_id or contract_id, not both.")
 
             query_header = self._make_request_header()
-            crypto_get_balance = crypto_get_account_balance_pb2.CryptoGetAccountBalanceQuery()
+            crypto_get_balance = (
+                crypto_get_account_balance_pb2.CryptoGetAccountBalanceQuery()
+            )
             crypto_get_balance.header.CopyFrom(query_header)
-            crypto_get_balance.accountID.CopyFrom(self.account_id._to_proto())
+
+            if self.account_id:
+                crypto_get_balance.accountID.CopyFrom(self.account_id._to_proto())
+            else:
+                crypto_get_balance.contractID.CopyFrom(self.contract_id._to_proto())
 
             query = query_pb2.Query()
-            if not hasattr(query, 'cryptogetAccountBalance'):
-                raise AttributeError("Query object has no attribute 'cryptogetAccountBalance'")
+            if not hasattr(query, "cryptogetAccountBalance"):
+                raise AttributeError(
+                    "Query object has no attribute 'cryptogetAccountBalance'"
+                )
             query.cryptogetAccountBalance.CopyFrom(crypto_get_balance)
 
             return query
@@ -73,7 +119,7 @@ def _make_request(self) -> query_pb2.Query:
     def _get_method(self, channel: _Channel) -> _Method:
         """
         Returns the appropriate gRPC method for the account balance query.
-        
+
         Implements the abstract method from Query to provide the specific
         gRPC method for getting account balances.
 
@@ -84,16 +130,15 @@ def _get_method(self, channel: _Channel) -> _Method:
             _Method: The method wrapper containing the query function
         """
         return _Method(
-            transaction_func=None,
-            query_func=channel.crypto.cryptoGetBalance
+            transaction_func=None, query_func=channel.crypto.cryptoGetBalance
         )
 
     def execute(self, client) -> AccountBalance:
         """
         Executes the account balance query.
-        
+
         This function delegates the core logic to `_execute()`, and may propagate exceptions raised by it.
-        
+
         Sends the query to the Hedera network and processes the response
         to return an AccountBalance object.
 
@@ -113,26 +158,28 @@ def execute(self, client) -> AccountBalance:
 
         return AccountBalance._from_proto(response.cryptogetAccountBalance)
 
-    def _get_query_response(self, response: Any) -> crypto_get_account_balance_pb2.CryptoGetAccountBalanceResponse:
+    def _get_query_response(
+        self, response: Any
+    ) -> crypto_get_account_balance_pb2.CryptoGetAccountBalanceResponse:
         """
         Extracts the account balance response from the full response.
-        
+
         Implements the abstract method from Query to extract the
         specific account balance response object.
-        
+
         Args:
             response: The full response from the network
-            
+
         Returns:
             The crypto get account balance response object
         """
         return response.cryptogetAccountBalance
-    
+
     def _is_payment_required(self):
         """
         Account balance query does not require payment.
-        
+
         Returns:
             bool: False
         """
-        return False
+        return False