fix: support public for batch_key in Transaction (hiero-ledger#1051)

AubreyDDD · web-flow · commit 3241d95fbeb8 · 2025-12-19T10:25:42.000Z
Signed-off-by: Aubrey Du &lt;duk110293@gmail.com&gt;
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -58,6 +58,7 @@ This changelog is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.
 - Move `account_allowance_delete_transaction_hbar.py` from `examples/` to `examples/account/` for better organization (#1003)
 - Improved consistency of transaction examples (#1120)
 - Refactored `account_create_transaction_with_fallback_alias.py` by splitting the monolithic `create_account_with_fallback_alias` function into modular functions: `generate_fallback_key`, `fetch_account_info`, and `print_account_summary`. The existing `setup_client()` function was reused for improved readability and structure (#1018)
+- Allow `PublicKey` for batch_key in `Transaction`, enabling both `PrivateKey` and `PublicKey` for batched transactions
 - Allow `PublicKey` for `TokenUpdateKeys` in `TokenUpdateTransaction`, enabling non-custodial workflows where operators can build transactions using only public keys (#934).
 - Bump protobuf toml to protobuf==6.33.2
 - chore: Move account allowance example to correct folder
diff --git a/examples/transaction/batch_transaction.py b/examples/transaction/batch_transaction.py
@@ -162,9 +162,9 @@ def transfer_token(client, sender, recipient, token_id):
 
 def perform_batch_tx(client, sender, recipient, token_id, freeze_key):
     """
-    Perform a batch transaction.
+    Perform a batch transaction using PrivateKey as batch_key.
     """
-    print("\nPerforming batch transaction (unfreeze → transfer → freeze)...")
+    print("\nPerforming batch transaction with PrivateKey (unfreeze → transfer → freeze)...")
     batch_key = PrivateKey.generate()
 
     unfreeze_tx = (
@@ -204,6 +204,58 @@ def perform_batch_tx(client, sender, recipient, token_id, freeze_key):
     print(f"Batch transaction status: {ResponseCode(receipt.status).name}")
 
 
+def perform_batch_tx_with_public_key(client, sender, recipient, token_id, freeze_key):
+    """
+    Perform a batch transaction using PublicKey as batch_key.
+    Demonstrates that batch_key can accept both PrivateKey and PublicKey.
+    """
+    print("\n✨ Performing batch transaction with PublicKey (unfreeze → transfer → freeze)...")
+    
+    # Generate a key pair - we'll use the PublicKey as batch_key
+    batch_private_key = PrivateKey.generate()
+    batch_public_key = batch_private_key.public_key()
+    
+    print(f"Using PublicKey as batch_key: {batch_public_key}")
+
+    # Create inner transactions using PublicKey as batch_key
+    unfreeze_tx = (
+        TokenUnfreezeTransaction()
+        .set_account_id(sender)
+        .set_token_id(token_id)
+        .batchify(client, batch_public_key)  # Using PublicKey!
+        .sign(freeze_key)
+    )
+
+    transfer_tx = (
+        TransferTransaction()
+        .add_token_transfer(token_id, sender, -1)
+        .add_token_transfer(token_id, recipient, 1)
+        .batchify(client, batch_public_key)  # Using PublicKey!
+    )
+
+    freeze_tx = (
+        TokenFreezeTransaction()
+        .set_account_id(sender)
+        .set_token_id(token_id)
+        .batchify(client, batch_public_key)  # Using PublicKey!
+        .sign(freeze_key)
+    )
+
+    # Assemble the batch transaction
+    batch = (
+        BatchTransaction()
+        .add_inner_transaction(unfreeze_tx)
+        .add_inner_transaction(transfer_tx)
+        .add_inner_transaction(freeze_tx)
+        .freeze_with(client)
+        .sign(batch_private_key)  # Sign with PrivateKey for execution
+    )
+
+    receipt = batch.execute(client)
+    print(f"Batch transaction with PublicKey status: {ResponseCode(receipt.status).name}")
+    print("   This demonstrates that batch_key now accepts both PrivateKey and PublicKey!")
+
+
 def main():
     client = setup_client()
     freeze_key = PrivateKey.generate()
@@ -227,19 +279,36 @@ def main():
     get_balance(client, client.operator_account_id, token_id)
     get_balance(client, recipient_id, token_id)
 
-    # Batch unfreeze → transfer → freeze
-    perform_batch_tx(
-        client, client.operator_account_id, recipient_id, token_id, freeze_key
-    )
+    # Batch unfreeze → transfer → freeze (using PrivateKey)
+    perform_batch_tx(client, client.operator_account_id, recipient_id, token_id, freeze_key)
 
-    print("\nBalances after batch:")
+    print("\nBalances after first batch:")
     get_balance(client, client.operator_account_id, token_id)
     get_balance(client, recipient_id, token_id)
-
-    # Should fail again Verify that token is again freeze for account
+    
+    # Verify that token is frozen again
+    receipt = transfer_token(client, client.operator_account_id, recipient_id, token_id)
+    if receipt.status == ResponseCode.ACCOUNT_FROZEN_FOR_TOKEN:
+        print("\n✅ Correct: Account is frozen again after first batch")
+    else:
+        print("\nAccount should be frozen again!")
+        sys.exit(1)
+    
+    # Now demonstrate using PublicKey as batch_key
+    print("\n" + "="*80)
+    print("Demonstrating PublicKey support for batch_key")
+    print("="*80)
+    
+    perform_batch_tx_with_public_key(client, client.operator_account_id, recipient_id, token_id, freeze_key)
+    
+    print("\nBalances after second batch (with PublicKey):")
+    get_balance(client, client.operator_account_id, token_id)
+    get_balance(client, recipient_id, token_id)
+    
+    # Verify that token is frozen again
     receipt = transfer_token(client, client.operator_account_id, recipient_id, token_id)
     if receipt.status == ResponseCode.ACCOUNT_FROZEN_FOR_TOKEN:
-        print("\nCorrect: Account is frozen again")
+        print("\n✅ Success! Account is frozen again, PublicKey batch_key works correctly!")
     else:
         print("\nAccount should be frozen again!")
         sys.exit(1)
diff --git a/src/hiero_sdk_python/transaction/transaction.py b/src/hiero_sdk_python/transaction/transaction.py
@@ -16,6 +16,7 @@
 from hiero_sdk_python.response_code import ResponseCode
 from hiero_sdk_python.transaction.transaction_id import TransactionId
 from hiero_sdk_python.transaction.transaction_response import TransactionResponse
+from hiero_sdk_python.utils.key_utils import Key, key_to_proto
 
 if TYPE_CHECKING:
     from hiero_sdk_python.schedule.schedule_create_transaction import (
@@ -65,7 +66,7 @@ def __init__(self) -> None:
         # changed from int: 2_000_000 to Hbar: 0.02
         self._default_transaction_fee = Hbar(0.02)
         self.operator_account_id = None  
-        self.batch_key: Optional[PrivateKey] = None
+        self.batch_key: Optional[Key] = None
 
     def _make_request(self):
         """
@@ -434,7 +435,7 @@ def build_base_transaction_body(self) -> transaction_pb2.TransactionBody:
         transaction_body.max_custom_fees.extend(custom_fee_limits)
 
         if self.batch_key:
-            transaction_body.batch_key.CopyFrom(self.batch_key.public_key()._to_proto())
+            transaction_body.batch_key.CopyFrom(key_to_proto(self.batch_key))
 
         return transaction_body
 
@@ -807,12 +808,12 @@ def _from_protobuf(cls, transaction_body, body_bytes: bytes, sig_map):
 
         return transaction
     
-    def set_batch_key(self, key: PrivateKey):
+    def set_batch_key(self, key: Key):
         """
         Set the batch key required for batch transaction.
 
         Args:
-            batch_key (PrivateKey): Private key to use as batch key.
+            batch_key (Key): Key to use as batch key (accepts both PrivateKey and PublicKey).
 
         Returns:
             Transaction: A reconstructed transaction instance of the appropriate subclass. 
@@ -821,13 +822,13 @@ def set_batch_key(self, key: PrivateKey):
         self.batch_key = key
         return self
     
-    def batchify(self, client: Client, batch_key: PrivateKey):
+    def batchify(self, client: Client, batch_key: Key):
         """
         Marks the current transaction as an inner (batched) transaction.
 
         Args:
             client (Client): The client instance to use for setting defaults.
-            batch_key (PrivateKey): Private key to use as batch key.
+            batch_key (Key): Key to use as batch key (accepts both PrivateKey and PublicKey).
         
         Returns:
             Transaction: A reconstructed transaction instance of the appropriate subclass.
diff --git a/tests/integration/batch_transaction_e2e_test.py b/tests/integration/batch_transaction_e2e_test.py
@@ -2,6 +2,7 @@
 import pytest
 from hiero_sdk_python.account.account_create_transaction import AccountCreateTransaction
 from hiero_sdk_python.crypto.private_key import PrivateKey
+from hiero_sdk_python.crypto.public_key import PublicKey
 from hiero_sdk_python.file.file_id import FileId
 from hiero_sdk_python.query.account_info_query import AccountInfoQuery
 from hiero_sdk_python.query.transaction_get_receipt_query import TransactionGetReceiptQuery
@@ -352,3 +353,154 @@ def test_batch_transaction_with_inner_schedule_transaction(env):
 
     batch_receipt = batch_tx.execute(env.client)
     assert batch_receipt.status == ResponseCode.BATCH_TRANSACTION_IN_BLACKLIST
+
+
+def test_batch_transaction_with_public_key_as_batch_key(env):
+    """Test batch transaction can use PublicKey as batch_key."""
+    # Generate a key pair - we'll use the PublicKey as batch_key
+    batch_private_key = PrivateKey.generate()
+    batch_public_key = batch_private_key.public_key()
+    
+    receiver_id = create_account_tx(PrivateKey.generate().public_key(), env.client)
+
+    # Use PublicKey in batchify
+    transfer_tx = (
+        TransferTransaction()
+        .add_hbar_transfer(account_id=env.operator_id, amount=-1)
+        .add_hbar_transfer(account_id=receiver_id, amount=1)
+        .batchify(env.client, batch_public_key)  # Using PublicKey!
+    )
+
+    # Verify batch_key was set to PublicKey
+    assert isinstance(transfer_tx.batch_key, PublicKey)
+    assert transfer_tx.batch_key == batch_public_key
+
+    # Sign and execute with PrivateKey
+    batch_tx = (
+        BatchTransaction()
+        .add_inner_transaction(transfer_tx)
+        .freeze_with(env.client)
+        .sign(batch_private_key)  # Sign with corresponding PrivateKey
+    )
+
+    batch_receipt = batch_tx.execute(env.client)
+    assert batch_receipt.status == ResponseCode.SUCCESS
+
+    # Inner Transaction Receipt
+    transfer_tx_id = batch_tx.get_inner_transaction_ids()[0]
+    transfer_tx_receipt = (
+        TransactionGetReceiptQuery()
+        .set_transaction_id(transfer_tx_id)
+        .execute(env.client)
+    )
+    assert transfer_tx_receipt.status == ResponseCode.SUCCESS
+
+
+def test_batch_transaction_with_mixed_public_and_private_keys(env):
+    """Test batch transaction can handle inner transactions with mixed PrivateKey and PublicKey."""
+    # Generate different keys
+    batch_key1_private = PrivateKey.generate()
+    batch_key2_private = PrivateKey.generate()
+    batch_key2_public = batch_key2_private.public_key()
+    batch_key3_private = PrivateKey.generate()
+    batch_key3_public = batch_key3_private.public_key()
+
+    # Create receivers
+    receiver_id1 = create_account_tx(PrivateKey.generate().public_key(), env.client)
+    receiver_id2 = create_account_tx(PrivateKey.generate().public_key(), env.client)
+    receiver_id3 = create_account_tx(PrivateKey.generate().public_key(), env.client)
+
+    # First inner transaction uses PrivateKey
+    transfer_tx1 = (
+        TransferTransaction()
+        .add_hbar_transfer(account_id=env.operator_id, amount=-1)
+        .add_hbar_transfer(account_id=receiver_id1, amount=1)
+        .batchify(env.client, batch_key1_private)  # PrivateKey
+    )
+
+    # Second inner transaction uses PublicKey
+    transfer_tx2 = (
+        TransferTransaction()
+        .add_hbar_transfer(account_id=env.operator_id, amount=-1)
+        .add_hbar_transfer(account_id=receiver_id2, amount=1)
+        .batchify(env.client, batch_key2_public)  # PublicKey
+    )
+
+    # Third inner transaction uses PublicKey
+    transfer_tx3 = (
+        TransferTransaction()
+        .add_hbar_transfer(account_id=env.operator_id, amount=-1)
+        .add_hbar_transfer(account_id=receiver_id3, amount=1)
+        .batchify(env.client, batch_key3_public)  # PublicKey
+    )
+
+    # Verify key types
+    assert isinstance(transfer_tx1.batch_key, PrivateKey)
+    assert isinstance(transfer_tx2.batch_key, PublicKey)
+    assert isinstance(transfer_tx3.batch_key, PublicKey)
+
+    # Assemble and sign batch transaction
+    batch_tx = (
+        BatchTransaction()
+        .add_inner_transaction(transfer_tx1)
+        .add_inner_transaction(transfer_tx2)
+        .add_inner_transaction(transfer_tx3)
+        .freeze_with(env.client)
+        .sign(batch_key1_private)
+        .sign(batch_key2_private)  # Sign with PrivateKey for PublicKey batch_key
+        .sign(batch_key3_private)  # Sign with PrivateKey for PublicKey batch_key
+    )
+
+    batch_receipt = batch_tx.execute(env.client)
+    assert batch_receipt.status == ResponseCode.SUCCESS
+
+    # Verify all inner transactions succeeded
+    for transfer_tx_id in batch_tx.get_inner_transaction_ids():
+        transfer_tx_receipt = (
+            TransactionGetReceiptQuery()
+            .set_transaction_id(transfer_tx_id)
+            .execute(env.client)
+        )
+        assert transfer_tx_receipt.status == ResponseCode.SUCCESS
+
+
+def test_batch_transaction_set_batch_key_with_public_key(env):
+    """Test batch transaction inner transaction can use set_batch_key with PublicKey."""
+    # Generate a key pair
+    batch_private_key = PrivateKey.generate()
+    batch_public_key = batch_private_key.public_key()
+    
+    receiver_id = create_account_tx(PrivateKey.generate().public_key(), env.client)
+
+    # Use set_batch_key with PublicKey instead of batchify
+    transfer_tx = (
+        TransferTransaction()
+        .add_hbar_transfer(account_id=env.operator_id, amount=-1)
+        .add_hbar_transfer(account_id=receiver_id, amount=1)
+        .set_batch_key(batch_public_key)  # Using set_batch_key with PublicKey
+        .freeze_with(env.client)
+        .sign(env.operator_key)  # Sign inner transaction with operator key
+    )
+
+    # Verify batch_key was set correctly
+    assert transfer_tx.batch_key == batch_public_key
+    assert isinstance(transfer_tx.batch_key, PublicKey)
+
+    batch_tx = (
+        BatchTransaction()
+        .add_inner_transaction(transfer_tx)
+        .freeze_with(env.client)
+        .sign(batch_private_key)  # Sign batch transaction with batch key
+    )
+
+    batch_receipt = batch_tx.execute(env.client)
+    assert batch_receipt.status == ResponseCode.SUCCESS
+
+    # Inner Transaction Receipt
+    transfer_tx_id = batch_tx.get_inner_transaction_ids()[0]
+    transfer_tx_receipt = (
+        TransactionGetReceiptQuery()
+        .set_transaction_id(transfer_tx_id)
+        .execute(env.client)
+    )
+    assert transfer_tx_receipt.status == ResponseCode.SUCCESS
diff --git a/tests/unit/batch_transaction_test.py b/tests/unit/batch_transaction_test.py
