feat: improve the integration test_case

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

Author: manishdait

src/hiero_sdk_python/consensus/topic_message.py

@@ -8,6 +8,7 @@
 
 from hiero_sdk_python.timestamp import Timestamp
 from hiero_sdk_python.hapi.mirror import consensus_service_pb2 as mirror_proto
+from hiero_sdk_python.transaction.transaction_id import TransactionId
 
 
 class TopicMessageChunk:
@@ -40,7 +41,7 @@ def __init__(
             consensus_timestamp: datetime,
             message_data: Dict[str, Union[bytes, int]],
             chunks: List[TopicMessageChunk],
-            transaction_id: Optional[str] = None,
+            transaction_id: Optional[TransactionId] = None,
     ) -> None:
         """
         Args:
@@ -52,14 +53,14 @@ def __init__(
                               "sequence_number": int
                           }
             chunks (List[TopicMessageChunk]): All individual chunks that form this message.
-            transaction_id (Optional[str]): The transaction ID string if available.
+            transaction_id (Optional[Transaction]): The transaction ID if available.
         """
         self.consensus_timestamp: datetime = consensus_timestamp
         self.contents: Union[bytes, int] = message_data["contents"]
         self.running_hash: Union[bytes, int] = message_data["running_hash"]
         self.sequence_number: Union[bytes, int] = message_data["sequence_number"]
         self.chunks: List[TopicMessageChunk] = chunks
-        self.transaction_id: Optional[str] = transaction_id
+        self.transaction_id: Optional[TransactionId] = transaction_id
 
     @classmethod
     def of_single(cls, response: mirror_proto.ConsensusTopicResponse) -> "TopicMessage":  # type: ignore
@@ -72,13 +73,9 @@ def of_single(cls, response: mirror_proto.ConsensusTopicResponse) -> "TopicMessa
         running_hash: Union[bytes, int] = response.runningHash
         sequence_number: Union[bytes, int] = chunk.sequence_number
 
-        transaction_id: Optional[str] = None
+        transaction_id: Optional[TransactionId] = None
         if response.HasField("chunkInfo") and response.chunkInfo.HasField("initialTransactionID"):
-            tx_id = response.chunkInfo.initialTransactionID
-            transaction_id = (
-                f"{tx_id.shardNum}.{tx_id.realmNum}.{tx_id.accountNum}-"
-                f"{tx_id.transactionValidStart.seconds}.{tx_id.transactionValidStart.nanos}"
-            )
+            transaction_id = TransactionId._from_proto(response.chunkInfo.initialTransactionID)
 
         return cls(
             consensus_timestamp,
@@ -102,7 +99,7 @@ def of_many(cls, responses: List[mirror_proto.ConsensusTopicResponse]) -> "Topic
 
         chunks: List[TopicMessageChunk] = []
         total_size: int = 0
-        transaction_id: Optional[str] = None
+        transaction_id: Optional[TransactionId] = None
 
         for r in sorted_responses:
             c = TopicMessageChunk(r)
@@ -116,12 +113,7 @@ def of_many(cls, responses: List[mirror_proto.ConsensusTopicResponse]) -> "Topic
                     and r.HasField("chunkInfo")
                     and r.chunkInfo.HasField("initialTransactionID")
             ):
-                tx_id = r.chunkInfo.initialTransactionID
-                transaction_id = (
-                    f"{tx_id.accountID.shardNum, '0'}.{tx_id.accountID.realmNum, '0'}.{tx_id.accountID.accountNum}-"
-                    f"{tx_id.transactionValidStart.seconds}.{tx_id.transactionValidStart.nanos}"
-                )
-                print(transaction_id)
+                transaction_id = TransactionId._from_proto(r.chunkInfo.initialTransactionID)
 
         contents = bytearray(total_size)
 

src/hiero_sdk_python/consensus/topic_message_submit_transaction.py

@@ -44,7 +44,8 @@ def __init__(
 
         self._current_index = 0
         self._total_chunks = self.get_required_chunks()
-        self._transaction_ids: List[TransactionId] = []
+        self._intial_transaction_id: Optional[TransactionId] = None
+        self._initial_transaction_id: List[TransactionId] = []
         self._signing_keys: List["PrivateKey"] = []
 
     def get_required_chunks(self) -> int:
@@ -96,13 +97,17 @@ def set_chunk_size(self, chunk_size: int) -> "TopicMessageSubmitTransaction":
         Set maximum chunk size in bytes.
         
         Args:
-            chunck_size (int):  The size of each chunk in bytes.
+            chunk_size (int):  The size of each chunk in bytes.
 
         Returns:
             TopicMessageSubmitTransaction: This transaction instance (for chaining).
         """
         self._require_not_frozen()
+        if chunk_size <= 0:
+            raise ValueError("chunk_size must be positive")
+        
         self.chunk_size = chunk_size
+        self._total_chunks = self.get_required_chunks()
         return self
 
     def set_max_chunks(self, max_chunks: int) -> "TopicMessageSubmitTransaction":
@@ -116,6 +121,9 @@ def set_max_chunks(self, max_chunks: int) -> "TopicMessageSubmitTransaction":
             TopicMessageSubmitTransaction: This transaction instance (for chaining).
         """
         self._require_not_frozen()
+        if max_chunks <= 0:
+            raise ValueError("max_chunks must be positive")
+        
         self.max_chunks = max_chunks
         return self
 
@@ -193,10 +201,10 @@ def _build_proto_body(self) -> consensus_submit_message_pb2.ConsensusSubmitMessa
             message=chunk_content
         )
 
-         # Multi-chunk metadata
+        # Multi-chunk metadata
         if self._total_chunks > 1:
             body.chunkInfo.CopyFrom(consensus_submit_message_pb2.ConsensusMessageChunkInfo(
-                initialTransactionID=self.transaction_id._to_proto(),
+                initialTransactionID=self._intial_transaction_id._to_proto(),
                 total=self._total_chunks,
                 number=self._current_index + 1
             ))
@@ -255,7 +263,9 @@ def freeze_with(self, client: "Client") -> "TopicMessageSubmitTransaction":
 
         for i in range(self.get_required_chunks()):
             if i == 0:
-                # First chunk uses the original transaction ID
+                if self._intial_transaction_id is None:
+                    self._intial_transaction_id = self.transaction_id
+
                 chunk_transaction_id = self.transaction_id
             else:
                 chunk_valid_start = timestamp_pb2.Timestamp(
@@ -268,15 +278,9 @@ def freeze_with(self, client: "Client") -> "TopicMessageSubmitTransaction":
                 )
 
             self._transaction_ids.append(chunk_transaction_id)
-    
-        for node in client.network.nodes:
-            self.node_account_id = node._account_id
-            transaction_body = self.build_transaction_body()
-            self._transaction_body_bytes[node._account_id] = transaction_body.SerializeToString()
-
-        self.node_account_id = client.network.current_node._account_id
+        
+        return super().freeze_with(client)
 
-        return self
 
     def execute(self, client: "Client"):
         self._validate_chunking()
@@ -286,14 +290,13 @@ def execute(self, client: "Client"):
 
         # Multi-chunk transaction - execute all chunks
         responses = []
-        print(self.get_required_chunks())
 
         for chunk_index in range(self.get_required_chunks()):
             self._current_index = chunk_index
 
             if self._transaction_ids and chunk_index < len(self._transaction_ids):
                 self.transaction_id = self._transaction_ids[chunk_index]
-            
+
             self._transaction_body_bytes.clear()
             self._signature_map.clear()
 

src/hiero_sdk_python/query/topic_message_query.py

@@ -1,12 +1,13 @@
 import time
 import threading
 from datetime import datetime
-from typing import Optional, Callable, Union, Dict, List, Any
+from typing import Optional, Callable, Union, Dict, List
 
 from hiero_sdk_python.hapi.mirror import consensus_service_pb2 as mirror_proto
 from hiero_sdk_python.hapi.services import basic_types_pb2, timestamp_pb2
 from hiero_sdk_python.consensus.topic_id import TopicId
 from hiero_sdk_python.consensus.topic_message import TopicMessage
+from hiero_sdk_python.transaction.transaction_id import TransactionId
 from hiero_sdk_python.utils.subscription_handle import SubscriptionHandle
 from hiero_sdk_python.client.client import Client
 
@@ -143,21 +144,15 @@ def run_stream():
                             on_message(msg_obj)
                             continue
 
-                        initial_tx_id = response.chunkInfo.initialTransactionID
-                        tx_id_str = (f"{initial_tx_id.accountID.shardNum or '0'}."
-                                     f"{initial_tx_id.accountID.realmNum or '0'}."
-                                     f"{initial_tx_id.accountID.accountNum}-"
-                                     f"{initial_tx_id.transactionValidStart.seconds}."
-                                     f"{initial_tx_id.transactionValidStart.nanos}")
+                        initial_tx_id = TransactionId._from_proto(response.chunkInfo.initialTransactionID)
 
-                        if tx_id_str not in pending_chunks:
-                            pending_chunks[tx_id_str] = []
+                        if initial_tx_id not in pending_chunks:
+                            pending_chunks[initial_tx_id] = []
 
-                        pending_chunks[tx_id_str].append(response)
+                        pending_chunks[initial_tx_id].append(response)
 
-                        if len(pending_chunks) == response.chunkInfo.total:
-                            chunk_list =  [resp for chunk_list in pending_chunks.values() for resp in chunk_list]
-                            pending_chunks.clear() 
+                        if len(pending_chunks[initial_tx_id]) == response.chunkInfo.total:
+                            chunk_list = pending_chunks.pop(initial_tx_id)
 
                             msg_obj = TopicMessage.of_many(chunk_list)
                             on_message(msg_obj)