chore: test

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

Author: manishdait

tests/unit/subscription_handle_test.py

@@ -6,17 +6,20 @@
 
 
 def test_not_cancelled_by_default():
+    """Test a new handle starts in a non-cancelled state."""
     handle = SubscriptionHandle()
     assert not handle.is_cancelled()
 
 
 def test_cancel_marks_as_cancelled():
+    """Test calling cancel updates the is_cancelled status."""
     handle = SubscriptionHandle()
     handle.cancel()
     assert handle.is_cancelled()
 
 
 def test_set_thread_and_join_calls_thread_join_with_timeout():
+    """Test that join correctly forwards the timeout to the underlying thread."""
     handle = SubscriptionHandle()
     mock_thread = Mock()
     handle.set_thread(mock_thread)
@@ -25,6 +28,25 @@ def test_set_thread_and_join_calls_thread_join_with_timeout():
 
 
 def test_join_without_thread_raises_nothing():
+    """Test join is a no-op if no thread has been associated."""
     handle = SubscriptionHandle()
     # should not raise
     handle.join()
+
+
+def test_cancel_triggers_grpc_termination():
+    """Test that cancelling the handle terminates the active gRPC call."""
+    handle = SubscriptionHandle()
+    mock_call = Mock()
+    handle._set_call(mock_call)
+    handle.cancel()
+    mock_call.cancel.assert_called_once()
+
+
+def test_immediate_cancellation_of_late_call():
+    """Test a gRPC call is cancelled immediately if set after the handle was cancelled."""
+    handle = SubscriptionHandle()
+    mock_call = Mock()
+    handle.cancel()
+    handle._set_call(mock_call)
+    mock_call.cancel.assert_called_once()

tests/unit/topic_message_query_test.py

@@ -1,15 +1,19 @@
 from __future__ import annotations
 
+import time
 from datetime import datetime, timezone
 from unittest.mock import MagicMock, patch
 
 import pytest
 
+from hiero_sdk_python.account.account_id import AccountId
 from hiero_sdk_python.client.client import Client
 from hiero_sdk_python.consensus.topic_id import TopicId
 from hiero_sdk_python.hapi.mirror import consensus_service_pb2 as mirror_proto
 from hiero_sdk_python.hapi.services import timestamp_pb2 as hapi_timestamp_pb2
+from hiero_sdk_python.hapi.services.consensus_submit_message_pb2 import ConsensusMessageChunkInfo
 from hiero_sdk_python.query.topic_message_query import TopicMessageQuery
+from hiero_sdk_python.transaction.transaction_id import TransactionId
 
 
 pytestmark = pytest.mark.unit
@@ -19,7 +23,8 @@
 def mock_client():
     """Fixture to provide a mock Client instance."""
     client = MagicMock(spec=Client)
-    client.operator_account_id = "0.0.12345"
+    client.operator_account_id = AccountId(0, 0, 12345)
+
     return client
 
 
@@ -40,11 +45,20 @@ def mock_subscription_response():
     )
 
 
+def test_topic_message_query_initialization():
+    """Test initializing the query with various parameter types and setters."""
+    start = datetime(2023, 1, 1, tzinfo=timezone.utc)
+    query = TopicMessageQuery().set_topic_id("0.0.123").set_start_time(start).set_limit(5).set_chunking_enabled(True)
+
+    assert query._topic_id.topicNum == 123
+    assert query._start_time.seconds == int(start.timestamp())
+    assert query._limit == 5
+    assert query._chunking_enabled is True
+
+
 # This test uses fixtures (mock_client, mock_topic_id, mock_subscription_response) as parameters
 def test_topic_message_query_subscription(mock_client, mock_topic_id, mock_subscription_response):
-    """
-    Test subscribing to topic messages using TopicMessageQuery.
-    """
+    """Test subscribing to topic messages using TopicMessageQuery."""
     query = TopicMessageQuery().set_topic_id(mock_topic_id).set_start_time(datetime.now(tz=timezone.utc))
 
     with patch("hiero_sdk_python.query.topic_message_query.TopicMessageQuery.subscribe") as mock_subscribe:
@@ -69,3 +83,29 @@ def side_effect(client, on_message, on_error):  # noqa: ARG001
         on_error.assert_not_called()
 
     print("Test passed: Subscription handled messages correctly.")
+
+
+def test_chunk_message_handling(mock_client):
+    """Test that multiple chunks are correctly buffered and released as a single message."""
+    query = TopicMessageQuery(topic_id="0.0.123", chunking_enabled=True)
+
+    # Mocking two chunks for the same transaction
+    tx_id = TransactionId.generate(mock_client.operator_account_id)._to_proto()
+    chunk1 = mirror_proto.ConsensusTopicResponse(
+        message=b"part1", chunkInfo=ConsensusMessageChunkInfo(initialTransactionID=tx_id, total=2, number=1)
+    )
+    chunk2 = mirror_proto.ConsensusTopicResponse(
+        message=b"part2", chunkInfo=ConsensusMessageChunkInfo(initialTransactionID=tx_id, total=2, number=2)
+    )
+
+    mock_client.mirror_stub.subscribeTopic.return_value = [chunk1, chunk2]
+
+    received_messages = []
+    query.subscribe(mock_client, on_message=lambda m: received_messages.append(m))
+
+    # Wait for thread execution
+    time.sleep(0.1)
+
+    assert len(received_messages) == 1
+    # Assuming TopicMessage.of_many joins messages
+    assert b"part1" in received_messages[0].message