feat(tck): Added getAccountBalance method to tck (hiero-ledger#2336)

Signed-off-by: Manish Dait <daitmanish88@gmail.com>

Author: manishdait

src/hiero_sdk_python/account/account_balance.py

@@ -16,18 +16,23 @@ class AccountBalance:
     Attributes:
         hbars (Hbar): The balance in hbars.
         token_balances (dict): A dictionary mapping TokenId to token balances.
+        token_decimals (dict, optional): A dictionary mapping TokenId to token deimals.
     """
 
-    def __init__(self, hbars: Hbar, token_balances: dict[TokenId, int] = None) -> None:
+    def __init__(
+        self, hbars: Hbar, token_balances: dict[TokenId, int] = None, token_decimals: dict[TokenId, int] = None
+    ) -> None:
         """
         Initializes the AccountBalance with the given hbar balance and token balances.
 
         Args:
             hbars (Hbar): The balance in hbars.
             token_balances (dict, optional): A dictionary mapping TokenId to token balances.
+            token_decimals (dict, optional): A dictionary mapping TokenId to token deimals.
         """
         self.hbars = hbars
         self.token_balances = token_balances or {}
+        self.token_decimals = token_decimals or {}
 
     @classmethod
     def _from_proto(cls, proto: CryptoGetAccountBalanceResponse) -> AccountBalance:
@@ -43,13 +48,18 @@ def _from_proto(cls, proto: CryptoGetAccountBalanceResponse) -> AccountBalance:
         hbars: Hbar = Hbar.from_tinybars(tinybars=proto.balance)
 
         token_balances: dict[TokenId, int] = {}
+        token_decimals: dict[TokenId, int] = {}
+
         if proto.tokenBalances:
             for token_balance in proto.tokenBalances:
                 token_id: TokenId = TokenId._from_proto(token_balance.tokenId)
                 balance: int = token_balance.balance
+                decimal: int = token_balance.decimals
+
                 token_balances[token_id] = balance
+                token_decimals[token_id] = decimal
 
-        return cls(hbars=hbars, token_balances=token_balances)
+        return cls(hbars=hbars, token_balances=token_balances, token_decimals=token_decimals)
 
     def __str__(self) -> str:
         """
@@ -59,10 +69,17 @@ def __str__(self) -> str:
             str: A string showing HBAR balance and token balances.
         """
         lines = [f"HBAR Balance: {self.hbars} hbars"]
+
         if self.token_balances:
             lines.append("Token Balances:")
             for token_id, balance in self.token_balances.items():
                 lines.append(f" - Token ID {token_id}: {balance} units")
+
+        if self.token_decimals:
+            lines.append("Token Decimals:")
+            for token_id, decimal in self.token_decimals.items():
+                lines.append(f" - Token ID {token_id}: {decimal} decimals")
+
         return "\n".join(lines)
 
     def __repr__(self) -> str:
@@ -77,4 +94,9 @@ def __repr__(self) -> str:
             if self.token_balances
             else "{}"
         )
-        return f"AccountBalance(hbars={self.hbars!r}, token_balances={token_balances_repr})"
+        token_decimals_repr = (
+            f"{{{', '.join(f'{token_id!r}: {decimals}' for token_id, decimals in self.token_decimals.items())}}}"
+            if self.token_decimals
+            else "{}"
+        )
+        return f"AccountBalance(hbars={self.hbars!r}, token_balances={token_balances_repr}) token_decimals={token_decimals_repr})"

src/hiero_sdk_python/query/account_balance_query.py

@@ -90,19 +90,17 @@ def _make_request(self) -> query_pb2.Query:
             Exception: If any other error occurs during request construction.
         """
         try:
-            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.header.CopyFrom(query_header)
 
-            if self.account_id:
+            if self.account_id is not None:
                 crypto_get_balance.accountID.CopyFrom(self.account_id._to_proto())
-            else:
+
+            if self.contract_id is not None:
                 crypto_get_balance.contractID.CopyFrom(self.contract_id._to_proto())
 
             query = query_pb2.Query()

tck/handlers/account.py

@@ -2,22 +2,32 @@
 
 from __future__ import annotations
 
+from hiero_sdk_python.account.account_balance import AccountBalance
 from hiero_sdk_python.account.account_create_transaction import AccountCreateTransaction
 from hiero_sdk_python.account.account_delete_transaction import AccountDeleteTransaction
 from hiero_sdk_python.account.account_id import AccountId
 from hiero_sdk_python.account.account_info import AccountInfo
 from hiero_sdk_python.account.account_update_transaction import AccountUpdateTransaction
+from hiero_sdk_python.contract.contract_id import ContractId
 from hiero_sdk_python.Duration import Duration
 from hiero_sdk_python.hbar import Hbar
+from hiero_sdk_python.query.account_balance_query import CryptoGetAccountBalanceQuery
 from hiero_sdk_python.query.account_info_query import AccountInfoQuery
 from hiero_sdk_python.response_code import ResponseCode
 from hiero_sdk_python.timestamp import Timestamp
 from hiero_sdk_python.transaction.transaction_receipt import TransactionReceipt
 from tck.handlers.registry import rpc_method
-from tck.param.account import CreateAccountParams, DeleteAccountParams, GetAccountInfoParams, UpdateAccountParams
+from tck.param.account import (
+    CreateAccountParams,
+    DeleteAccountParams,
+    GetAccountBalanceParams,
+    GetAccountInfoParams,
+    UpdateAccountParams,
+)
 from tck.response.account import (
     CreateAccountResponse,
     DeleteAccountResponse,
+    GetAccountBalanceResponse,
     GetAccountInfoResponse,
     StakingInfoResponse,
     TokenRelationshipResponse,
@@ -254,3 +264,30 @@ def delete_account(params: DeleteAccountParams) -> DeleteAccountResponse:
     receipt: TransactionReceipt = response.get_receipt(client, validate_status=True)
 
     return DeleteAccountResponse(status=ResponseCode(receipt.status).name)
+
+
+@rpc_method("getAccountBalance")
+def get_account_balance(params: GetAccountBalanceParams) -> GetAccountBalanceResponse:
+    """Get account balance for an account."""
+    client = get_client(params.sessionId)
+
+    query = CryptoGetAccountBalanceQuery().set_grpc_deadline(DEFAULT_GRPC_TIMEOUT)
+
+    if params.accountId is not None:
+        query.set_account_id(AccountId.from_string(params.accountId))
+    if params.contractId is not None:
+        query.set_contract_id(ContractId.from_string(params.contractId))
+
+    account_balance = query.execute(client)
+    return map_account_balance_response(account_balance)
+
+
+def map_account_balance_response(account_balance: AccountBalance) -> GetAccountBalanceResponse:
+    """Map AccountBalance class to the GetAccountBalanceResponse."""
+    token_balances = {str(token_id): int(balance) for token_id, balance in account_balance.token_balances.items()}
+
+    token_decimals = {str(token_id): int(decimals) for token_id, decimals in account_balance.token_decimals.items()}
+
+    return GetAccountBalanceResponse(
+        hbars=str(account_balance.hbars.to_tinybars()), tokenBalances=token_balances, tokenDecimals=token_decimals
+    )

tck/param/account.py

@@ -113,3 +113,18 @@ class GetAccountInfoParams(BaseParams):
     def parse_json_params(cls, params: dict) -> GetAccountInfoParams:
         """Parse JSON-RPC params into a GetAccountInfoParams instance."""
         return cls(accountId=params.get("accountId"), sessionId=parse_session_id(params))
+
+
+@dataclass
+class GetAccountBalanceParams(BaseParams):
+    """Request parameters for the getAccountBalance endpoint."""
+
+    accountId: str | None = None
+    contractId: str | None = None
+
+    @classmethod
+    def parse_json_params(cls, params):
+        """Parse JSON-RPC params into a GetAccountBalanceParams instance."""
+        return cls(
+            accountId=params.get("accountId"), contractId=params.get("contractId"), sessionId=parse_session_id(params)
+        )

tck/response/account.py

@@ -80,3 +80,12 @@ class GetAccountInfoResponse:
     tokenAllowances: list[dict] = field(default_factory=list)
     nftAllowances: list[dict] = field(default_factory=list)
     stakingInfo: StakingInfoResponse | None = None
+
+
+@dataclass
+class GetAccountBalanceResponse:
+    """Response payload for getAccountBalance."""
+
+    hbars: str | None = None
+    tokenBalances: dict[str, int] = field(default_factory=dict)
+    tokenDecimals: dict[str, int] = field(default_factory=dict)

tests/integration/account_balance_query_e2e_test.py

@@ -6,6 +6,7 @@
 from hiero_sdk_python.contract.contract_create_transaction import ContractCreateTransaction
 from hiero_sdk_python.contract.contract_delete_transaction import ContractDeleteTransaction
 from hiero_sdk_python.contract.contract_id import ContractId
+from hiero_sdk_python.exceptions import PrecheckError
 from hiero_sdk_python.query.account_balance_query import CryptoGetAccountBalanceQuery
 from hiero_sdk_python.response_code import ResponseCode
 from tests.integration.utils import IntegrationTestEnv
@@ -89,11 +90,10 @@ def test_integration_contract_balance_query_can_execute():
 def test_integration_balance_query_raises_when_neither_source_set():
     env = IntegrationTestEnv()
     try:
-        with pytest.raises(
-            ValueError,
-            match=r"Either account_id or contract_id must be set before making the request\.",
-        ):
+        with pytest.raises(PrecheckError) as err:
             CryptoGetAccountBalanceQuery().execute(env.client)
+
+        assert err.value.status == ResponseCode.INVALID_ACCOUNT_ID
     finally:
         env.close()
 

tests/unit/account_balance_query_test.py

@@ -132,15 +132,13 @@ def test_last_wins_when_both_account_id_and_contract_id_are_set(
     assert query.account_id is None
 
 
-def test_make_request_raises_when_neither_account_id_nor_contract_id_is_set():
-    """_make_request should raise if neither account_id nor contract_id is set."""
+def test_make_request_when_neither_account_id_nor_contract_id_is_set():
+    """_make_request should create request without ids when neither account_id nor contract_id is set."""
     query = CryptoGetAccountBalanceQuery()
+    proto = query._make_request()
 
-    with pytest.raises(
-        ValueError,
-        match=r"Either account_id or contract_id must be set before making the request\.",
-    ):
-        query._make_request()
+    assert not proto.cryptogetAccountBalance.HasField("accountID")
+    assert not proto.cryptogetAccountBalance.HasField("contractID")
 
 
 def test_make_request_populates_contract_id_only():

tests/unit/account_balance_test.py

@@ -5,6 +5,9 @@
 import pytest
 
 from hiero_sdk_python.account.account_balance import AccountBalance
+from hiero_sdk_python.account.account_id import AccountId
+from hiero_sdk_python.hapi.services.basic_types_pb2 import TokenBalance
+from hiero_sdk_python.hapi.services.crypto_get_account_balance_pb2 import CryptoGetAccountBalanceResponse
 from hiero_sdk_python.hbar import Hbar
 from hiero_sdk_python.tokens.token_id import TokenId
 
@@ -26,13 +29,14 @@ def test_account_balance_str_with_hbars_only():
     assert "Token Balances:" not in result
 
 
-def test_account_balance_str_with_token_balances():
-    """Test __str__ method with hbars and token balances."""
+def test_account_balance_str_with_token_balances_and_decimal():
+    """Test __str__ method with hbars and token balances and decimals."""
     hbars = Hbar(10)
     token_id_1 = TokenId(0, 0, 100)
     token_id_2 = TokenId(0, 0, 200)
     token_balances = {token_id_1: 1000, token_id_2: 500}
-    account_balance = AccountBalance(hbars=hbars, token_balances=token_balances)
+    token_decimals = {token_id_1: 1, token_id_2: 2}
+    account_balance = AccountBalance(hbars=hbars, token_balances=token_balances, token_decimals=token_decimals)
 
     result = str(account_balance)
 
@@ -42,10 +46,13 @@ def test_account_balance_str_with_token_balances():
     assert "Token Balances:" in result
     assert " - Token ID 0.0.100: 1000 units" in result
     assert " - Token ID 0.0.200: 500 units" in result
+    assert "Token Decimals:" in result
+    assert " - Token ID 0.0.100: 1 decimals" in result
+    assert " - Token ID 0.0.200: 2 decimals" in result
 
 
-def test_account_balance_str_with_empty_token_balances():
-    """Test __str__ method with empty token balances dict."""
+def test_account_balance_str_with_empty_token_balances_and_decimals():
+    """Test __str__ method with empty token balances and decimals dict."""
     hbars = Hbar(5.5)
     account_balance = AccountBalance(hbars=hbars, token_balances={})
 
@@ -56,6 +63,7 @@ def test_account_balance_str_with_empty_token_balances():
     assert " hbars" in result
     # Should not include token balances section when empty
     assert "Token Balances:" not in result
+    assert "Token Decimals:" not in result
 
 
 def test_account_balance_repr_with_hbars_only():
@@ -77,7 +85,8 @@ def test_account_balance_repr_with_token_balances():
     token_id_1 = TokenId(0, 0, 100)
     token_id_2 = TokenId(0, 0, 200)
     token_balances = {token_id_1: 1000, token_id_2: 500}
-    account_balance = AccountBalance(hbars=hbars, token_balances=token_balances)
+    token_decimals = {token_id_1: 1, token_id_2: 2}
+    account_balance = AccountBalance(hbars=hbars, token_balances=token_balances, token_decimals=token_decimals)
 
     result = repr(account_balance)
 
@@ -87,3 +96,28 @@ def test_account_balance_repr_with_token_balances():
     assert "0.0.100" in result or "TokenId" in result
     assert "1000" in result
     assert "500" in result
+    assert "token_decimals=" in result
+    assert "1" in result
+    assert "2" in result
+
+
+def test_create_account_balance_from_proto():
+    token_blances_proto = [
+        TokenBalance(tokenId=TokenId(0, 0, 100)._to_proto(), balance=100, decimals=1),
+        TokenBalance(tokenId=TokenId(0, 0, 102)._to_proto(), balance=0, decimals=0),
+    ]
+
+    proto = CryptoGetAccountBalanceResponse(
+        accountID=AccountId(0, 0, 1)._to_proto(), balance=10, tokenBalances=token_blances_proto
+    )
+
+    account_balance = AccountBalance._from_proto(proto=proto)
+
+    assert account_balance is not None
+    assert account_balance.hbars.to_tinybars() == 10
+
+    assert len(account_balance.token_balances) == 2
+    assert account_balance.token_balances == {TokenId(0, 0, 100): 100, TokenId(0, 0, 102): 0}
+
+    assert len(account_balance.token_decimals) == 2
+    assert account_balance.token_decimals == {TokenId(0, 0, 100): 1, TokenId(0, 0, 102): 0}