feat: expand functionality token associate (#846)

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

Author: AntonioCeppellini

CHANGELOG.md

@@ -35,12 +35,14 @@ This changelog is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.
 - `staked_account_id`, `staked_node_id` and `decline_staking_reward` fields to AccountInfo
 - Added `examples/token_create_transaction_supply_key.py` to demonstrate token creation with and without a supply key.
 - Added `examples/token_create_transaction_kyc_key.py` to demonstrate KYC key functionality, including creating tokens with/without KYC keys, granting/revoking KYC status, and understanding KYC requirements for token transfers.
+- Add `set_token_ids`, `_from_proto`, `_validate_checksum` to TokenAssociateTransaction [#795]
 - Added BatchTransaction class
 - Add support for token metadata (bytes, max 100 bytes) in `TokenCreateTransaction`, including a new `set_metadata` setter, example, and tests. [#799]
 - Added `examples/token_create_transaction_token_fee_schedule.py` to demonstrate creating tokens with custom fee schedules and the consequences of not having it.
 - Added `examples/token_create_transaction_wipe_key.py` to demonstrate token wiping and the role of the wipe key.
 - Added `examples/account_allowance_approve_transaction_hbar.py` and `examples/account_allowance_delete_transaction_hbar.py`, deleted `examples/account_allowance_hbar.py`. [#775]
 - Added `docs\sdk_developers\training\receipts.md` as a training guide for users to understand hedera receipts.
+- Add `set_token_ids`, `_from_proto`, `_validate_checksum` to TokenAssociateTransaction [#795]
 - docs: added `network_and_client.md` with a table of contents, and added external example scripts (`client.py`).
 
 ### Changed

examples/tokens/token_associate_transaction.py

@@ -54,7 +54,6 @@ def setup_client():
         print("❌ Error: Please check OPERATOR_ID and OPERATOR_KEY in your .env file.")
         sys.exit(1)
 
-
 def create_test_account(client, operator_key):
     """
     Create a new test account for demonstration.
@@ -158,7 +157,7 @@ def associate_token_with_account(client, token_id, account_id, account_key):
             .sign(account_key)
             .execute(client)
         )
-    
+
         if receipt.status != ResponseCode.SUCCESS:
             print(
                 f"❌ Token association failed with status: {ResponseCode(receipt.status).name}"
@@ -170,6 +169,83 @@ def associate_token_with_account(client, token_id, account_id, account_key):
         print(f"❌ Error associating token with account: {e}")
         sys.exit(1)
 
+
+def associate_two_tokens_mixed_types_with_set_token_ids(client, token_id_1, token_id_2, account_id, account_key):
+    """
+    Associate two tokens using set_token_ids() with mixed types:
+    - first as TokenId
+    - second as string
+    """
+    try:
+        receipt = (
+            TokenAssociateTransaction()
+            .set_account_id(account_id)
+            .set_token_ids(
+                [
+                    token_id_1,  # TokenId instance
+                    str(token_id_2),  # string representation → converted internally
+                ]
+            )
+            .freeze_with(client)
+            .sign(account_key)
+            .execute(client)
+        )
+
+        if receipt.status != ResponseCode.SUCCESS:
+            print(
+                f"❌ Token association (mixed types) failed with status: "
+                f"{ResponseCode(receipt.status).name}"
+            )
+            sys.exit(1)
+
+        print("✅ Success! Token association completed.")
+        print(
+            f"   Account {account_id} can now hold and transfer tokens {token_id_1} and {token_id_2}"
+        )
+    except Exception as e:
+        print(f"❌ Error in while associating tokens: {e}")
+        sys.exit(1)
+
+
+def demonstrate_invalid_set_token_ids_usage(client, account_id, account_key):
+    """
+    Example 4: demonstrate that set_token_ids() rejects invalid types,
+    i.e. values that are neither TokenId nor string.
+    """
+    print("`set_token_ids()` only accepts a list of TokenId or strings (also mixed)")
+    invalid_value = 123  # ❌ This type is not supported from `set_token_ids()`
+    print(f"Trying to associate a token using a list of {type(invalid_value)}")
+
+    try:
+        receipt = (
+            TokenAssociateTransaction()
+            .set_account_id(account_id)
+            .set_token_ids([invalid_value])  # this should fail
+            .freeze_with(client)
+            .sign(account_key)
+            .execute(client)
+        )
+
+        if receipt.status != ResponseCode.SUCCESS:
+            print(
+                f"❌ Token creation failed with status: {ResponseCode(receipt.status).name}"
+            )
+            sys.exit(1)
+
+    except Exception as e:
+        if (
+            type(e) == TypeError
+            and "Invalid token_id type: expected TokenId or str, got" in e.args[0]
+        ):
+            print(
+                "✅ Correct behavior: invalid token_id type was rejected from `set_token_ids`"
+            )
+            print(f"   Error: {e}")
+        else:
+            print(f"❌ Unexpected error while creating transaction: {e}")
+            sys.exit(1)
+
+
 def verify_token_association(client, account_id, token_id):
     """
     Verify that a token is properly associated with an account.
@@ -231,23 +307,43 @@ def main():
     print("\nSTEP 2: Creating a new account...")
     account_id, account_private_key = create_test_account(client, operator_key)
 
-    # Step 3: Create a new token
-    print("\nSTEP 3: Creating a new fungible token...")
-    token_id = create_fungible_token(client, operator_id, operator_key)
-
-    # Step 4: Associate the token with the new account
-    print(f"\nSTEP 4: Associating token {token_id} with account {account_id}...")
-    associate_token_with_account(client, token_id, account_id, account_private_key)
-
-    # Step 5: Verify the token association
-    print(f"\nSTEP 5: Verifying token association...")
-    is_associated = verify_token_association(client, account_id, token_id)
-
+    # step 3: How to not use set_token_ids
+    print("\nSTEP 3: Demonstrating invalid input handling in set_token_ids...")
+    demonstrate_invalid_set_token_ids_usage(client, account_id, account_private_key)
+
+    # Step 4: new tokens
+    print("\nSTEP 4: Creating new fungible tokens...")
+    token_id_0 = create_fungible_token(client, operator_id, operator_key)
+    token_id_1 = create_fungible_token(client, operator_id, operator_key)
+    token_id_2 = create_fungible_token(client, operator_id, operator_key)
+
+    # Step 5: Associate a single token with the new account
+    print(f"\nSTEP 5: Associating token {token_id_0} with account {account_id}...")
+    associate_token_with_account(client, token_id_0, account_id, account_private_key)
+
+    # Step 6: Associate multiple tokens with the new account
+    print(
+        f"\nSTEP 6: Associating token {token_id_1} and token {token_id_2} with account {account_id}..."
+    )
+    associate_two_tokens_mixed_types_with_set_token_ids(
+        client, token_id_1, token_id_2, account_id, account_private_key
+    )
+
+    # Step 7: Verify the token association
+    print(f"\nSTEP 7: Verifying token association...")
+    is_associated = verify_token_association(client, account_id, token_id_0)
+    is_associated = verify_token_association(client, account_id, token_id_1)
+    is_associated = verify_token_association(client, account_id, token_id_2)
+
+    tokens = [token_id_0, token_id_1, token_id_2]
     # Summary
     print("\n" + "=" * 50)
     print("🎉 Token Association Demo Completed Successfully!")
     print(f"   New Account: {account_id}")
-    print(f"   New Token: {token_id}")
+    print("   New Tokens:")
+
+    for token in tokens:
+        print(f"    -{token}")
     print(f"   Association: {'✅ VERIFIED' if is_associated else '❌ FAILED'}")
     print("=" * 50)
 

src/hiero_sdk_python/tokens/token_associate_transaction.py

@@ -5,26 +5,28 @@
 Provides TokenAssociateTransaction, a subclass of Transaction for associating
 tokens with accounts on the Hedera network using the Hedera Token Service (HTS) API.
 """
-from typing import Optional, List
+
+from typing import Optional, List, Union
 
 from hiero_sdk_python.account.account_id import AccountId
 from hiero_sdk_python.channels import _Channel
 from hiero_sdk_python.executable import _Method
 from hiero_sdk_python.hapi.services import token_associate_pb2, transaction_pb2
-from hiero_sdk_python.hapi.services.schedulable_transaction_body_pb2 import (
-    SchedulableTransactionBody,
-)
+from hiero_sdk_python.hapi.services.schedulable_transaction_body_pb2 import SchedulableTransactionBody
+
 from hiero_sdk_python.tokens.token_id import TokenId
 from hiero_sdk_python.transaction.transaction import Transaction
 
+TokenIdLike = Union[TokenId, str]
+
 
 class TokenAssociateTransaction(Transaction):
     """
     Represents a token associate transaction on the Hedera network.
 
     This transaction associates the specified tokens with an account,
     allowing the account to hold and transact with those tokens.
-    
+
     Inherits from the base Transaction class and implements the required methods
     to build and execute a token association transaction.
     """
@@ -47,7 +49,7 @@ def __init__(
         self._default_transaction_fee: int = 500_000_000
 
     def set_account_id(self, account_id: AccountId) -> "TokenAssociateTransaction":
-        """        
+        """
         Sets the account ID for the token association transaction.
         Args:
             account_id (AccountId): The account ID to associate tokens with.
@@ -64,13 +66,41 @@ def add_token_id(self, token_id: TokenId) -> "TokenAssociateTransaction":
         self.token_ids.append(token_id)
         return self
 
+    def set_token_ids(self, token_ids: List[TokenId]) -> "TokenAssociateTransaction":
+        """
+        Sets the list of token IDs for the token association transaction.
+
+        This mirrors the JavaScript SDK's `setTokenIds()` API,
+        providing a convenient way to associate multiple tokens at once.
+
+        Args:
+            token_ids: Iterable of TokenId instances or string representations.
+
+        Returns:
+            TokenAssociateTransaction.
+        """
+        self._require_not_frozen()
+        tokens_to_add: List[TokenId] = []
+        for token_id in token_ids:
+            if isinstance(token_id, TokenId):
+                tokens_to_add.append(token_id)
+            elif isinstance(token_id, str):
+                tokens_to_add.append(TokenId.from_string(token_id))
+            else:
+                raise TypeError(
+                    f"Invalid token_id type: expected TokenId or str, got {type(token_id).__name__}"
+                )
+
+        self.token_ids = tokens_to_add
+        return self
+
     def _build_proto_body(self) -> token_associate_pb2.TokenAssociateTransactionBody:
         """
         Returns the protobuf body for the token associate transaction.
-        
+
         Returns:
             TokenAssociateTransactionBody: The protobuf body for this transaction.
-            
+
         Raises:
             ValueError: If account ID or token IDs are not set.
         """
@@ -81,7 +111,26 @@ def _build_proto_body(self) -> token_associate_pb2.TokenAssociateTransactionBody
             account=self.account_id._to_proto(),
             tokens=[token_id._to_proto() for token_id in self.token_ids]
         )
-        
+
+    @classmethod
+    def _from_proto(
+        cls, body: token_associate_pb2.TokenAssociateTransactionBody
+    ) -> "TokenAssociateTransaction":
+        """
+        Construct a TokenAssociateTransaction from its protobuf.
+        """
+        account_id = AccountId._from_proto(body.account)
+        token_ids: List[TokenId] = []
+
+        for proto_token in body.tokens:
+            token_id = TokenId._from_proto(proto_token)
+            token_ids.append(token_id)
+
+        return cls(
+            account_id=account_id,
+            token_ids=token_ids,
+        )
+
     def build_transaction_body(self) -> transaction_pb2.TransactionBody:
         """
         Builds and returns the protobuf transaction body for token association.
@@ -94,7 +143,7 @@ def build_transaction_body(self) -> transaction_pb2.TransactionBody:
         transaction_body.tokenAssociate.CopyFrom(token_associate_body)
 
         return transaction_body
-        
+
     def build_scheduled_body(self) -> SchedulableTransactionBody:
         """
         Builds the scheduled transaction body for this token associate transaction.
@@ -107,6 +156,19 @@ def build_scheduled_body(self) -> SchedulableTransactionBody:
         schedulable_body.tokenAssociate.CopyFrom(token_associate_body)
         return schedulable_body
 
+    def _validate_checksums(self, client) -> None:
+        """
+        Validate the checksums for all IDs (account + tokens) in this transaction.
+
+        Mirrors the style used across the SDK: each ID calls its own
+        validate_checksum(client) method.
+        """
+        if self.account_id is not None:
+            self.account_id.validate_checksum(client)
+
+        for token_id in self.token_ids:
+            token_id.validate_checksum(client)
+
     def _get_method(self, channel: _Channel) -> _Method:
         return _Method(
             transaction_func=channel.token.associateTokens,

tests/integration/token_associate_transaction_e2e_test.py

@@ -33,12 +33,25 @@ def test_integration_token_associate_transaction_can_execute():
         token_id = create_fungible_token(env)
         assert token_id is not None, "TokenID not found in receipt. Token may not have been created."
 
-        associate_transaction = TokenAssociateTransaction(
-            account_id=new_account_id,
-            token_ids=[token_id]
+        associate_transaction = (
+            TokenAssociateTransaction()
+            .set_account_id(new_account_id)
+            .set_token_ids([token_id])
         )
-        
+
         associate_transaction.freeze_with(env.client)
+
+        # building transaction body
+        transaction_body = associate_transaction.build_transaction_body()
+
+        proto_body = transaction_body.tokenAssociate
+
+        # Recreating transaction from protobuf
+        transaction_from_proto = TokenAssociateTransaction._from_proto(proto_body)
+
+        assert transaction_from_proto.account_id == associate_transaction.account_id
+        assert transaction_from_proto.token_ids == associate_transaction.token_ids
+
         associate_transaction.sign(new_account_private_key)
         associate_receipt = associate_transaction.execute(env.client)
 
@@ -55,4 +68,4 @@ def test_integration_token_associate_transaction_can_execute():
 
         assert dissociate_receipt.status == ResponseCode.SUCCESS, f"Token dissociation failed with status: {ResponseCode(dissociate_receipt.status).name}"
     finally:
-        env.close() 
+        env.close()