fix: INVALID_NODE_ACCOUNT during node switching (#93)

* refactor: improve transaction handling and signing

Signed-off-by: dosi <[email protected]>

* refactor: update query to build correctly payment transaction

Signed-off-by: dosi <[email protected]>

* test: fix unit tests

Signed-off-by: dosi <[email protected]>

* refactor: update _require_frozen to properly check if dict is empty

Signed-off-by: dosi <[email protected]>

* refactor: make transaction_body_bytes and signature_map private in Transaction class

Signed-off-by: dosi <[email protected]>

* docs: add comments to Transaction class

Signed-off-by: dosi <[email protected]>

* refactor: update query and query_payment to use the right Transaction attributes

Signed-off-by: dosi <[email protected]>

* test: fix unit tests

Signed-off-by: dosi <[email protected]>

* refactor: update Transaction class to address PR comments

Signed-off-by: dosi <[email protected]>

* revert: remove unintended changes from test_executable.py

Signed-off-by: dosi <[email protected]>

---------

Signed-off-by: dosi <[email protected]>

Author: Dosik13

src/hiero_sdk_python/query/query.py

@@ -145,7 +145,7 @@ def _build_query_payment_transaction(self, payer_account_id, payer_private_key,
         tx.transaction_id = TransactionId.generate(payer_account_id)
 
         body_bytes = tx.build_transaction_body().SerializeToString()
-        tx.transaction_body_bytes = body_bytes
+        tx._transaction_body_bytes.setdefault(node_account_id, body_bytes)
         tx.sign(payer_private_key)
 
         return tx.to_proto()

src/hiero_sdk_python/transaction/query_payment.py

@@ -26,7 +26,7 @@ def build_query_payment_transaction(
     tx.transaction_id = TransactionId.generate(payer_account_id)
 
     body_bytes = tx.build_transaction_body().SerializeToString()
-    tx.transaction_body_bytes = body_bytes
+    tx._transaction_body_bytes.setdefault(node_account_id, body_bytes)
 
     tx.sign(payer_private_key)
     return tx.to_proto()

src/hiero_sdk_python/transaction/transaction.py

@@ -1,5 +1,6 @@
 import hashlib
 
+from hiero_sdk_python.account.account_id import AccountId
 from hiero_sdk_python.exceptions import PrecheckError
 from hiero_sdk_python.executable import _Executable, _ExecutionState
 from hiero_sdk_python.hapi.services import (basic_types_pb2, transaction_body_pb2, transaction_contents_pb2, transaction_pb2)
@@ -8,7 +9,6 @@
 from hiero_sdk_python.transaction.transaction_id import TransactionId
 from hiero_sdk_python.transaction.transaction_response import TransactionResponse
 
-
 class Transaction(_Executable):
     """
     Base class for all Hedera transactions.
@@ -34,8 +34,17 @@ def __init__(self):
         self.transaction_valid_duration = 120 
         self.generate_record = False
         self.memo = ""
-        self.transaction_body_bytes = None
-        self.signature_map = basic_types_pb2.SignatureMap()
+        # Maps each node's AccountId to its corresponding transaction body bytes
+        # This allows us to maintain separate transaction bodies for each node
+        # which is necessary in case node is unhealthy and we have to switch it with other node.
+        # Each transaction body has the AccountId of the node it's being submitted to.
+        # If these do not match `INVALID_NODE_ACCOUNT` error will occur.
+        self._transaction_body_bytes: dict[AccountId, bytes] = {}
+        
+        # Maps transaction body bytes to their associated signatures
+        # This allows us to maintain the signatures for each unique transaction
+        # and ensures that the correct signatures are used when submitting transactions
+        self._signature_map: dict[bytes, basic_types_pb2.SignatureMap] = {}
         self._default_transaction_fee = 2_000_000
         self.operator_account_id = None  
 
@@ -148,18 +157,24 @@ def sign(self, private_key):
         """
         # We require the transaction to be frozen before signing
         self._require_frozen()
+        
+        # We sign the bodies for each node in case we need to switch nodes during execution.
+        for body_bytes in self._transaction_body_bytes.values():
+            signature = private_key.sign(body_bytes)
 
-        signature = private_key.sign(self.transaction_body_bytes)
-
-        public_key_bytes = private_key.public_key().to_bytes_raw()
+            public_key_bytes = private_key.public_key().to_bytes_raw()
 
-        sig_pair = basic_types_pb2.SignaturePair(
-            pubKeyPrefix=public_key_bytes,
-            ed25519=signature
-        )
+            sig_pair = basic_types_pb2.SignaturePair(
+                pubKeyPrefix=public_key_bytes,
+                ed25519=signature
+            )
 
-        self.signature_map.sigPair.append(sig_pair)
+            # We initialize the signature map for this body_bytes if it doesn't exist yet
+            self._signature_map.setdefault(body_bytes, basic_types_pb2.SignatureMap())
 
+            # Append the signature pair to the signature map for this transaction body
+            self._signature_map[body_bytes].sigPair.append(sig_pair)
+        
         return self
 
     def to_proto(self):
@@ -175,9 +190,17 @@ def to_proto(self):
         # We require the transaction to be frozen before converting to protobuf
         self._require_frozen()
 
+        body_bytes = self._transaction_body_bytes.get(self.node_account_id)
+        if body_bytes is None:
+            raise ValueError(f"No transaction body found for node {self.node_account_id}")
+
+        sig_map = self._signature_map.get(body_bytes)
+        if sig_map is None:
+            raise ValueError("No signature map found for the current transaction body")
+
         signed_transaction = transaction_contents_pb2.SignedTransaction(
-            bodyBytes=self.transaction_body_bytes,
-            sigMap=self.signature_map
+            bodyBytes=body_bytes,
+            sigMap=sig_map
         )
 
         return transaction_pb2.Transaction(
@@ -197,18 +220,22 @@ def freeze_with(self, client):
         Raises:
             Exception: If required IDs are not set.
         """
-        if self.transaction_body_bytes is not None:
+        if self._transaction_body_bytes:
             return self
-
+        
         if self.transaction_id is None:
             self.transaction_id = client.generate_transaction_id()
-
-        if self.node_account_id is None:
-            self.node_account_id = client.network.current_node._account_id
-
-        # print(f"Transaction's node account ID set to: {self.node_account_id}")
-        self.transaction_body_bytes = self.build_transaction_body().SerializeToString()
-
+        
+        # We iterate through every node in the client's network
+        # For each node, set the node_account_id and build the transaction body
+        # This allows the transaction to be submitted to any node in the network
+        for node in client.network.nodes:
+            self.node_account_id = node._account_id
+            self._transaction_body_bytes[node._account_id] = self.build_transaction_body().SerializeToString()
+        
+        # Set the node account id to the current node in the network
+        self.node_account_id = client.network.current_node._account_id
+        
         return self
 
     def execute(self, client):
@@ -228,7 +255,7 @@ def execute(self, client):
             MaxAttemptsError: If the transaction/query fails after the maximum number of attempts
             ReceiptStatusError: If the query fails with a receipt status error
         """
-        if self.transaction_body_bytes is None:
+        if not self._transaction_body_bytes:
             self.freeze_with(client)
 
         if self.operator_account_id is None:
@@ -257,8 +284,13 @@ def is_signed_by(self, public_key):
             bool: True if signed by the given public key, False otherwise.
         """
         public_key_bytes = public_key.to_bytes_raw()
-
-        for sig_pair in self.signature_map.sigPair:
+        
+        sig_map = self._signature_map.get(self._transaction_body_bytes.get(self.node_account_id))
+        
+        if sig_map is None:
+            return False
+        
+        for sig_pair in sig_map.sigPair:
             if sig_pair.pubKeyPrefix == public_key_bytes:
                 return True
         return False
@@ -317,7 +349,7 @@ def _require_not_frozen(self):
         Raises:
             Exception: If the transaction has already been frozen.
         """
-        if self.transaction_body_bytes is not None:
+        if self._transaction_body_bytes:
             raise Exception("Transaction is immutable; it has been frozen.")
 
     def _require_frozen(self):
@@ -330,7 +362,7 @@ def _require_frozen(self):
         Raises:
             Exception: If the transaction has not been frozen yet.
         """
-        if self.transaction_body_bytes is None:
+        if not self._transaction_body_bytes:
             raise Exception("Transaction is not frozen")
 
     def set_transaction_memo(self, memo):

tests/unit/test_account_create_transaction.py

@@ -55,8 +55,8 @@ def test_account_create_transaction_build(mock_account_ids):
     assert transaction_body.cryptoCreateAccount.initialBalance == 100000000
     assert transaction_body.cryptoCreateAccount.memo == "Test account"
 
-# This test uses fixture mock_account_ids as parameter
-def test_account_create_transaction_sign(mock_account_ids):
+# This test uses fixture (mock_account_ids, mock_client) as parameter
+def test_account_create_transaction_sign(mock_account_ids, mock_client):
     """Test signing the account create transaction."""
     operator_id, node_account_id = mock_account_ids
 
@@ -72,12 +72,23 @@ def test_account_create_transaction_sign(mock_account_ids):
     )
     account_tx.transaction_id = generate_transaction_id(operator_id)
     account_tx.node_account_id = node_account_id
-    account_tx.freeze_with(None)  
-    account_tx.sign(operator_private_key)
+    account_tx.freeze_with(mock_client)
 
-    # Verify signature was added
-    assert len(account_tx.signature_map.sigPair) == 1, \
+    # Add first signiture
+    account_tx.sign(mock_client.operator_private_key)
+    body_bytes = account_tx._transaction_body_bytes[node_account_id]
+    
+    assert body_bytes in account_tx._signature_map, "Body bytes should be a key in the signature map dictionary"
+    assert len(account_tx._signature_map[body_bytes].sigPair) == 1, \
         "Transaction should have exactly one signature"
+    
+    # Add second signiture
+    account_tx.sign(operator_private_key)
+    body_bytes = account_tx._transaction_body_bytes[node_account_id]
+    
+    assert body_bytes in account_tx._signature_map, "Body bytes should be a key in the signature map dictionary"
+    assert len(account_tx._signature_map[body_bytes].sigPair) == 2, \
+        "Transaction should have exactly two signatures"
 
 def test_account_create_transaction():
     """Integration test for AccountCreateTransaction with retry and response handling."""

tests/unit/test_executable.py

@@ -380,3 +380,51 @@ def test_topic_create_transaction_fails_on_nonretriable_error():
 
         # Verify the error contains the expected status
         assert str(ResponseCode.INVALID_TRANSACTION_BODY) in str(exc_info.value)
+        
+def test_transaction_node_switching_body_bytes():
+    """Test that execution switches nodes after receiving a non-retriable error."""
+    ok_response = TransactionResponseProto(nodeTransactionPrecheckCode=ResponseCode.OK)
+    error = RealRpcError(grpc.StatusCode.UNAVAILABLE, "Test error")
+    
+    receipt_response = response_pb2.Response(
+        transactionGetReceipt=transaction_get_receipt_pb2.TransactionGetReceiptResponse(
+            header=response_header_pb2.ResponseHeader(
+                nodeTransactionPrecheckCode=ResponseCode.OK
+            ),
+            receipt=transaction_receipt_pb2.TransactionReceipt(
+                status=ResponseCode.SUCCESS
+            )
+        )
+    )
+    # First node gives error, second node gives OK, third node gives error
+    response_sequences = [
+        [error],
+        [ok_response, receipt_response],
+    ]
+    
+    with mock_hedera_servers(response_sequences) as client, patch('time.sleep'):
+        # We set the current node to 0
+        client.network._node_index = 0
+        client.network.current_node = client.network.nodes[0]
+        
+        transaction = (
+            AccountCreateTransaction()
+            .set_key(PrivateKey.generate().public_key())
+            .set_initial_balance(100_000_000)
+            .freeze_with(client)
+            .sign(client.operator_private_key)
+        )
+        
+        for node in client.network.nodes:
+            assert transaction._transaction_body_bytes.get(node._account_id) is not None, "Transaction body bytes should be set for all nodes"
+            sig_map = transaction._signature_map.get(transaction._transaction_body_bytes[node._account_id])
+            assert sig_map is not None, "Signature map should be set for all nodes"
+            assert len(sig_map.sigPair) == 1, "Signature map should have one signature"
+            assert sig_map.sigPair[0].pubKeyPrefix == client.operator_private_key.public_key().to_bytes_raw(), "Signature should be for the operator"
+
+        try:
+            transaction.execute(client)
+        except (Exception, grpc.RpcError) as e:
+            pytest.fail(f"Transaction execution should not raise an exception, but raised: {e}")
+        # Verify we're now on the second node
+        assert client.network.current_node._account_id == AccountId(0, 0, 4), "Client should have switched to the second node"

tests/unit/test_token_associate_transaction.py

@@ -69,9 +69,13 @@ def test_sign_transaction(mock_account_ids, mock_client):
 
     # Sign the transaction
     associate_tx.sign(private_key)
+    
+    node_id = mock_client.network.current_node._account_id
+    body_bytes = associate_tx._transaction_body_bytes[node_id]
 
-    assert len(associate_tx.signature_map.sigPair) == 1
-    sig_pair = associate_tx.signature_map.sigPair[0]
+    assert body_bytes in associate_tx._signature_map, "Body bytes should be a key in the signature map dictionary"
+    assert len(associate_tx._signature_map[body_bytes].sigPair) == 1
+    sig_pair = associate_tx._signature_map[body_bytes].sigPair[0]
 
     assert sig_pair.pubKeyPrefix == b'public_key'  
     assert sig_pair.ed25519 == b'signature'

tests/unit/test_token_create_transaction.py

@@ -276,20 +276,23 @@ def test_sign_transaction(mock_account_ids, mock_client):
     # Sign with both sign keys
     token_tx.sign(private_key) # Necessary
     token_tx.sign(private_key_admin) # Since admin key exists
+    
+    node_id = mock_client.network.current_node._account_id
+    body_bytes = token_tx._transaction_body_bytes[node_id]
 
     # Expect 2 sigPairs
-    assert len(token_tx.signature_map.sigPair) == 2
+    assert len(token_tx._signature_map[body_bytes].sigPair) == 2
 
-    sig_pair = token_tx.signature_map.sigPair[0]
+    sig_pair = token_tx._signature_map[body_bytes].sigPair[0]
     assert sig_pair.pubKeyPrefix == b"public_key"
     assert sig_pair.ed25519 == b"signature"
 
-    sig_pair_admin = token_tx.signature_map.sigPair[1]
+    sig_pair_admin = token_tx._signature_map[body_bytes].sigPair[1]
     assert sig_pair_admin.pubKeyPrefix == b"admin_public_key"
     assert sig_pair_admin.ed25519 == b"admin_signature"
 
     # Confirm that neither sigPair belongs to supply_key or freeze_key:
-    for sig_pair in token_tx.signature_map.sigPair:
+    for sig_pair in token_tx._signature_map[body_bytes].sigPair:
         assert sig_pair.pubKeyPrefix not in (b"supply_public_key", b"freeze_public_key")
 
 # This test uses fixture (mock_account_ids, mock_client) as parameter
@@ -445,7 +448,7 @@ def test_transaction_execution_failure(mock_account_ids):
     token_tx.transaction_id = generate_transaction_id(treasury_account)
 
     # Set the transaction body bytes to avoid calling build_transaction_body
-    token_tx.transaction_body_bytes = b"mock_body_bytes"
+    token_tx._transaction_body_bytes = b"mock_body_bytes"
 
     # Mock the client and its operator_private_key
     token_tx.client = MagicMock()
@@ -543,8 +546,8 @@ def test_overwrite_defaults(mock_account_ids, mock_client):
     # Confirm no adminKey was set
     assert not tx_body.tokenCreation.HasField("adminKey")
 
-# This test uses fixture mock_account_ids as parameter
-def test_transaction_freeze_prevents_modification(mock_account_ids):
+# This test uses fixture (mock_account_ids, mock_client) as parameter
+def test_transaction_freeze_prevents_modification(mock_account_ids, mock_client):
     """
     Test that after freeze() is called, attempts to modify TokenCreateTransaction
     parameters raise an exception indicating immutability.
@@ -562,10 +565,9 @@ def test_transaction_freeze_prevents_modification(mock_account_ids):
 
     transaction.node_account_id = node_account_id
     transaction.transaction_id = generate_transaction_id(treasury_account)
-    transaction.client = MagicMock()
 
     # Freeze the transaction
-    transaction.freeze_with(transaction.client)
+    transaction.freeze_with(mock_client)
 
     # Attempt to overwrite after freeze - expect exceptions
     with pytest.raises(Exception, match="Transaction is immutable; it has been frozen."):

tests/unit/test_token_delete_transaction.py

@@ -60,8 +60,11 @@ def test_sign_transaction(mock_account_ids, mock_client):
 
     delete_tx.sign(private_key)
 
-    assert len(delete_tx.signature_map.sigPair) == 1
-    sig_pair = delete_tx.signature_map.sigPair[0]
+    node_id = mock_client.network.current_node._account_id
+    body_bytes = delete_tx._transaction_body_bytes[node_id]
+
+    assert len(delete_tx._signature_map[body_bytes].sigPair) == 1
+    sig_pair = delete_tx._signature_map[body_bytes].sigPair[0]
     assert sig_pair.pubKeyPrefix == b'public_key'
     assert sig_pair.ed25519 == b'signature'
 

tests/unit/test_token_dissociate_transaction.py

@@ -89,8 +89,11 @@ def test_sign_transaction(mock_account_ids, mock_client):
 
     dissociate_tx.sign(private_key)
 
-    assert len(dissociate_tx.signature_map.sigPair) == 1
-    sig_pair = dissociate_tx.signature_map.sigPair[0]
+    node_id = mock_client.network.current_node._account_id
+    body_bytes = dissociate_tx._transaction_body_bytes[node_id]
+
+    assert len(dissociate_tx._signature_map[body_bytes].sigPair) == 1
+    sig_pair = dissociate_tx._signature_map[body_bytes].sigPair[0]
 
     assert sig_pair.pubKeyPrefix == b'public_key'  
     assert sig_pair.ed25519 == b'signature'

tests/unit/test_token_freeze_transaction.py

@@ -79,8 +79,11 @@ def test_sign_transaction(mock_account_ids, mock_client):
 
     freeze_tx.sign(freeze_key)
 
-    assert len(freeze_tx.signature_map.sigPair) == 1
-    sig_pair = freeze_tx.signature_map.sigPair[0]
+    node_id = mock_client.network.current_node._account_id
+    body_bytes = freeze_tx._transaction_body_bytes[node_id]
+
+    assert len(freeze_tx._signature_map[body_bytes].sigPair) == 1
+    sig_pair = freeze_tx._signature_map[body_bytes].sigPair[0]
     assert sig_pair.pubKeyPrefix == b'public_key'
     assert sig_pair.ed25519 == b'signature'