fix: correct token association check (hiero-ledger#1044)

Signed-off-by: Aubrey Du <[email protected]>

Author: AubreyDDD

CHANGELOG.md

@@ -59,6 +59,7 @@ This changelog is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.
 
 ### Fixed
 
+- Fix token association verification in `token_airdrop_transaction.py` to correctly check if tokens are associated by using `token_id in token_balances` instead of incorrectly displaying zero balances which was misleading (#[815])
 - Fixed inactivity bot workflow not checking out repository before running (#964)
 - Fixed the topic_message_query integarion test
 - good first issue template yaml rendering

examples/tokens/token_airdrop_transaction.py

@@ -17,13 +17,13 @@
  TokenAirdropTransaction,
  TokenAssociateTransaction,
  TokenMintTransaction,
- CryptoGetAccountBalanceQuery,
  TokenType,
  ResponseCode,
  NftId,
  TransactionRecordQuery,
  TokenNftInfoQuery
 )
+from hiero_sdk_python.query.account_info_query import AccountInfoQuery
 
 load_dotenv()
 network_name = os.getenv('NETWORK', 'testnet').lower()
@@ -139,14 +139,16 @@ def associate_tokens(client, recipient_id, recipient_key, tokens):
         assocciate_tx.sign(recipient_key)
         assocciate_tx.execute(client)
 
-        balance_before = (
-            CryptoGetAccountBalanceQuery(account_id=recipient_id)
-            .execute(client)
-            .token_balances
-        )
-        print("Tokens associated with recipient:")
-        print(f"{tokens[0]}: {balance_before.get(tokens[0])}")
-        print(f"{tokens[1]}: {balance_before.get(tokens[1])}")
+        # Use AccountInfoQuery to check token associations
+        info = AccountInfoQuery(account_id=recipient_id).execute(client)
+        
+        # Get the list of associated token IDs from token_relationships
+        associated_token_ids = [rel.token_id for rel in info.token_relationships]
+
+        print("\nTokens associated with recipient:")
+        for token_id in tokens:
+            associated = token_id in associated_token_ids
+            print(f"  {token_id}: {'Associated' if associated else 'NOT Associated'}")
 
         print("\n✅ Success! Token association complete.")
 
@@ -282,6 +284,7 @@ def verify_post_airdrop_balances(
 ):
     """
     Verify post-airdrop balances and NFT ownership to confirm successful transfer.
+    Uses AccountInfoQuery for robust token association and balance verification.
     Accepts a `balances_before` mapping with keys 'sender' and 'recipient'.
     Returns (fully_verified_bool, details_dict)
     """
@@ -298,18 +301,20 @@ def verify_post_airdrop_balances(
         "record_checks": record_verification,
     }
 
-    # Get current balances (fail-fast is not desired; capture errors in details)
+    # Get current account info and balances using AccountInfoQuery (more robust)
     try:
-        sender_current = CryptoGetAccountBalanceQuery(account_id=operator_id).execute(client).token_balances
+        sender_info = AccountInfoQuery(account_id=operator_id).execute(client)
+        sender_current = {rel.token_id: rel.balance for rel in sender_info.token_relationships}
     except Exception as e:
-        details["error"] = f"Error fetching sender balance: {e}"
+        details["error"] = f"Error fetching sender account info: {e}"
         details["fully_verified"] = False
         return False, details
 
     try:
-        recipient_current = CryptoGetAccountBalanceQuery(account_id=recipient_id).execute(client).token_balances
+        recipient_info = AccountInfoQuery(account_id=recipient_id).execute(client)
+        recipient_current = {rel.token_id: rel.balance for rel in recipient_info.token_relationships}
     except Exception as e:
-        details["error"] = f"Error fetching recipient balance: {e}"
+        details["error"] = f"Error fetching recipient account info: {e}"
         details["fully_verified"] = False
         return False, details
 
@@ -349,16 +354,24 @@ def verify_post_airdrop_balances(
 
 
 def _log_balances_before(client, operator_id, recipient_id, token_id, nft_id):
-    """Fetch and print sender/recipient token balances before airdrop."""
-    print("\nStep 5: Checking balances before airdrop...")
-    sender_balances_before = CryptoGetAccountBalanceQuery(account_id=operator_id).execute(client).token_balances
-    recipient_balances_before = CryptoGetAccountBalanceQuery(account_id=recipient_id).execute(client).token_balances
+    """Fetch and print sender/recipient token balances and associations before airdrop."""
+    print("\nStep 5: Checking balances and associations before airdrop...")
+    
+    # Get account info to verify associations
+    sender_info = AccountInfoQuery(account_id=operator_id).execute(client)
+    recipient_info = AccountInfoQuery(account_id=recipient_id).execute(client)
+    
+    # Build dictionaries for balances from token_relationships (more robust than balance query)
+    sender_balances_before = {rel.token_id: rel.balance for rel in sender_info.token_relationships}
+    recipient_balances_before = {rel.token_id: rel.balance for rel in recipient_info.token_relationships}
+    
     print(f"Sender ({operator_id}) balances before airdrop:")
     print(f"  {token_id}: {sender_balances_before.get(token_id, 0)}")
     print(f"  {nft_id}: {sender_balances_before.get(nft_id, 0)}")
     print(f"Recipient ({recipient_id}) balances before airdrop:")
     print(f"  {token_id}: {recipient_balances_before.get(token_id, 0)}")
     print(f"  {nft_id}: {recipient_balances_before.get(nft_id, 0)}")
+    
     return sender_balances_before, recipient_balances_before
 
 
@@ -377,8 +390,18 @@ def _print_verification_summary(record_result, verification_details, operator_id
 
 def _print_final_summary(client, operator_id, recipient_id, token_id, nft_id, serial_number, verification_details):
     """Print final balances after airdrop and a small summary table."""
-    sender_balances_after = verification_details.get("sender_balances_after") or CryptoGetAccountBalanceQuery(account_id=operator_id).execute(client).token_balances
-    recipient_balances_after = verification_details.get("recipient_balances_after") or CryptoGetAccountBalanceQuery(account_id=recipient_id).execute(client).token_balances
+    # Use AccountInfoQuery for accurate balance retrieval if not already cached
+    if verification_details.get("sender_balances_after") is None:
+        sender_info = AccountInfoQuery(account_id=operator_id).execute(client)
+        sender_balances_after = {rel.token_id: rel.balance for rel in sender_info.token_relationships}
+    else:
+        sender_balances_after = verification_details.get("sender_balances_after")
+    
+    if verification_details.get("recipient_balances_after") is None:
+        recipient_info = AccountInfoQuery(account_id=recipient_id).execute(client)
+        recipient_balances_after = {rel.token_id: rel.balance for rel in recipient_info.token_relationships}
+    else:
+        recipient_balances_after = verification_details.get("recipient_balances_after")
 
     print("\nBalances after airdrop:")
     print(f"Sender ({operator_id}):")