chore: added test for the changes

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

Author: manishdait

src/hiero_sdk_python/query/topic_message_query.py

@@ -52,7 +52,7 @@ def __init__(
         self._topic_id: basic_types_pb2.TopicID | None = self._parse_topic_id(topic_id) if topic_id else None
         self._start_time: timestamp_pb2.Timestamp | None = self._parse_timestamp(start_time) if start_time else None
         self._end_time: timestamp_pb2.Timestamp | None = self._parse_timestamp(end_time) if end_time else None
-        self._limit: int | None = limit
+        self._limit: int = limit if limit is not None else 0
         self._chunking_enabled: bool = chunking_enabled
 
         self._max_attempts: int = 10
@@ -147,6 +147,59 @@ def _parse_timestamp(self, dt: datetime) -> timestamp_pb2.Timestamp:
         nanos = int((dt.timestamp() - seconds) * 1e9)
         return timestamp_pb2.Timestamp(seconds=seconds, nanos=nanos)
 
+    def _build_query_request(self, state: SubscriptionState) -> mirror_proto.ConsensusTopicQuery:
+        """Build the request object based on current subscription state."""
+        request = mirror_proto.ConsensusTopicQuery(topicID=self._topic_id)
+
+        if self._end_time is not None:
+            request.consensusEndTime.CopyFrom(self._end_time)
+
+        if state.last_message is not None:
+            last_message_time = state.last_message.consensusTimestamp
+
+            seconds = last_message_time.seconds
+            nanos = last_message_time.nanos + 1
+
+            if nanos >= 1_000_000_000:
+                seconds += 1
+                nanos = 0
+
+            request.consensusStartTime.seconds = seconds
+            request.consensusStartTime.nanos = nanos
+
+            if self._limit > 0:
+                request.limit = max(0, self._limit - state.count)
+        else:
+            if self._start_time is not None:
+                request.consensusStartTime.CopyFrom(self._start_time)
+            request.limit = self._limit
+
+        return request
+
+    def _handle_response(self, response, state: SubscriptionState, on_message: Callable[[TopicMessage], None]) -> None:
+        """Handles single or chunked messages."""
+        state.count += 1
+        state.last_message = response
+
+        if not self._chunking_enabled or not response.HasField("chunkInfo") or response.chunkInfo.total <= 1:
+            message = TopicMessage.of_single(response)
+            on_message(message)
+            return
+
+        initial_tx_id = TransactionId._from_proto(response.chunkInfo.initialTransactionID)
+
+        if initial_tx_id not in state.pending_messages:
+            state.pending_messages[initial_tx_id] = []
+
+        chunks = state.pending_messages[initial_tx_id]
+        chunks.append(response)
+
+        if len(chunks) == response.chunkInfo.total:
+            del state.pending_messages[initial_tx_id]
+
+            message = TopicMessage.of_many(chunks)
+            on_message(message)
+
     def subscribe(
         self,
         client: Client,
@@ -164,22 +217,7 @@ def subscribe(
 
         def run_stream():
             while state.attempt < self._max_attempts and not subscription_handle.is_cancelled():
-                request = mirror_proto.ConsensusTopicQuery(topicID=self._topic_id)
-
-                if self._end_time is not None:
-                    request.consensusEndTime.CopyFrom(self._end_time)
-
-                if state.last_message is not None:
-                    last_message_time = state.last_message.consensusTimestamp
-                    request.consensusStartTime.seconds = last_message_time.seconds
-                    request.consensusStartTime.nanos = last_message_time.nanos + 1
-
-                    if self._limit > 0:
-                        request.limit = max(0, self._limit - state.count)
-                else:
-                    if self._start_time is not None:
-                        request.consensusStartTime.CopyFrom(self._start_time)
-                    request.limit = self._limit
+                request = self._build_query_request(state)
 
                 try:
                     message_stream = client.mirror_stub.subscribeTopic(request)
@@ -189,31 +227,7 @@ def run_stream():
                         if subscription_handle.is_cancelled():
                             return
 
-                        state.count += 1
-                        state.last_message = response
-
-                        if (
-                            not self._chunking_enabled
-                            or not response.HasField("chunkInfo")
-                            or response.chunkInfo.total <= 1
-                        ):
-                            message = TopicMessage.of_single(response)
-                            on_message(message)
-                            continue
-
-                        initial_tx_id = TransactionId._from_proto(response.chunkInfo.initialTransactionID)
-
-                        if initial_tx_id not in state.pending_messages:
-                            state.pending_messages[initial_tx_id] = []
-
-                        chunks = state.pending_messages[initial_tx_id]
-                        chunks.append(response)
-
-                        if len(chunks) == response.chunkInfo.total:
-                            del state.pending_messages[initial_tx_id]
-
-                            message = TopicMessage.of_many(chunks)
-                            on_message(message)
+                        self._handle_response(response, state, on_message)
 
                     if self._completion_handler:
                         self._completion_handler()

tests/unit/topic_message_query_test.py

@@ -4,6 +4,7 @@
 from datetime import datetime, timezone
 from unittest.mock import MagicMock, patch
 
+import grpc
 import pytest
 
 from hiero_sdk_python.account.account_id import AccountId
@@ -14,6 +15,7 @@
 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
+from tests.unit.mock_server import RealRpcError
 
 
 pytestmark = pytest.mark.unit
@@ -24,6 +26,7 @@ def mock_client():
     """Fixture to provide a mock Client instance."""
     client = MagicMock(spec=Client)
     client.operator_account_id = AccountId(0, 0, 12345)
+    client.mirror_stub = MagicMock()
 
     return client
 
@@ -39,7 +42,7 @@ def mock_subscription_response():
     """Fixture to provide a mock response from a topic subscription."""
     return mirror_proto.ConsensusTopicResponse(
         consensusTimestamp=hapi_timestamp_pb2.Timestamp(seconds=12345, nanos=67890),
-        message=b"Hello, world!",
+        message=b"Hello Hiero!",
         runningHash=b"\x00" * 48,
         sequenceNumber=1,
     )
@@ -56,7 +59,49 @@ def test_topic_message_query_initialization():
     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_invalid_max_backoff():
+    """Test that invalid max_backoff raises errors."""
+    query = TopicMessageQuery()
+
+    with pytest.raises(ValueError, match="max_backoff must be at least 500 ms"):
+        query.set_max_backoff(0.1)
+
+
+def test_topic_message_query_invalid_max_attempts():
+    """Test that invalid max_attempts raises errors."""
+    query = TopicMessageQuery()
+
+    with pytest.raises(ValueError, match="max_attempts must be greater than 0"):
+        query.set_max_attempts(0)
+
+
+def test_topic_message_query_invalid_topic_id():
+    """Test that invalid topic_id raises errors."""
+    query = TopicMessageQuery()
+
+    # Invalid TopicId type
+    with pytest.raises(TypeError, match="Invalid topic_id format"):
+        query.set_topic_id(12345)
+
+    # Invalid TopicId format
+    with pytest.raises(ValueError, match="Invalid topic ID string"):
+        query.set_topic_id("12345")
+
+
+def test_subscribe_missing_config(mock_client):
+    """Test that subscribe fails if Topic ID or Mirror Stub is missing."""
+    # No TopicId
+    query_no_id = TopicMessageQuery()
+    with pytest.raises(ValueError, match="Topic ID must be set before subscribing"):
+        query_no_id.subscribe(mock_client, on_message=MagicMock())
+
+    # No MirrorStub
+    query_ok = TopicMessageQuery(topic_id="0.0.123")
+    mock_client.mirror_stub = None
+    with pytest.raises(ValueError, match="Client has no mirror_stub"):
+        query_ok.subscribe(mock_client, on_message=MagicMock())
+
+
 def test_topic_message_query_subscription(mock_client, mock_topic_id, mock_subscription_response):
     """Test subscribing to topic messages using TopicMessageQuery."""
     query = TopicMessageQuery().set_topic_id(mock_topic_id).set_start_time(datetime.now(tz=timezone.utc))
@@ -77,35 +122,110 @@ def side_effect(client, on_message, on_error):  # noqa: ARG001
         called_args = on_message.call_args[0][0]
         assert called_args.consensusTimestamp.seconds == 12345
         assert called_args.consensusTimestamp.nanos == 67890
-        assert called_args.message == b"Hello, world!"
+        assert called_args.message == b"Hello Hiero!"
         assert called_args.sequenceNumber == 1
 
         on_error.assert_not_called()
 
-    print("Test passed: Subscription handled messages correctly.")
 
-
-def test_chunk_message_handling(mock_client):
+def test_chunk_message_handling_improved(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()
+    tx_id = TransactionId.generate(mock_client.operator_account_id)
+    tx_id_proto = tx_id._to_proto()
+
     chunk1 = mirror_proto.ConsensusTopicResponse(
-        message=b"part1", chunkInfo=ConsensusMessageChunkInfo(initialTransactionID=tx_id, total=2, number=1)
+        message=b"chunk-1", chunkInfo=ConsensusMessageChunkInfo(initialTransactionID=tx_id_proto, total=2, number=1)
     )
     chunk2 = mirror_proto.ConsensusTopicResponse(
-        message=b"part2", chunkInfo=ConsensusMessageChunkInfo(initialTransactionID=tx_id, total=2, number=2)
+        message=b"chunk-2", chunkInfo=ConsensusMessageChunkInfo(initialTransactionID=tx_id_proto, total=2, number=2)
     )
 
-    mock_client.mirror_stub.subscribeTopic.return_value = [chunk1, chunk2]
+    mock_client.mirror_stub.subscribeTopic.return_value = iter([chunk1, chunk2])
 
     received_messages = []
-    query.subscribe(mock_client, on_message=lambda m: received_messages.append(m))
+    handle = query.subscribe(mock_client, on_message=lambda m: received_messages.append(m))
 
-    # Wait for thread execution
-    time.sleep(0.1)
+    handle._thread.join(timeout=1.0)
 
     assert len(received_messages) == 1
-    # Assuming TopicMessage.of_many joins messages
-    assert b"part1" in received_messages[0].message
+    assert b"chunk-1" in received_messages[0].contents
+    assert b"chunk-2" in received_messages[0].contents
+
+
+@pytest.mark.parametrize(
+    "error",
+    [
+        RealRpcError(grpc.StatusCode.NOT_FOUND, "unavailable"),
+        RealRpcError(grpc.StatusCode.UNAVAILABLE, "unavailable"),
+        RealRpcError(grpc.StatusCode.RESOURCE_EXHAUSTED, "busy"),
+        RealRpcError(grpc.StatusCode.INTERNAL, "received rst stream"),  # internal with rst stream
+        Exception("non grpc exception"),  # non grpc exception
+    ],
+)
+def test_retry_logic_on_retryable_error(mock_client, error):
+    """Test that the query retries on retryable errors but stops after max_attempts."""
+    query = TopicMessageQuery(topic_id="0.0.123").set_max_attempts(2).set_max_backoff(0.5)
+
+    mock_client.mirror_stub.subscribeTopic.side_effect = [error, error]
+
+    handle = query.subscribe(mock_client, on_message=MagicMock(), on_error=MagicMock())
+
+    handle._thread.join(timeout=2.0)
+
+    assert mock_client.mirror_stub.subscribeTopic.call_count == 2
+
+
+@pytest.mark.parametrize(
+    "non_retryable_error",
+    [
+        RealRpcError(grpc.StatusCode.PERMISSION_DENIED, "permission denied"),
+        RealRpcError(grpc.StatusCode.INVALID_ARGUMENT, "invalid argument"),
+        RealRpcError(grpc.StatusCode.UNAUTHENTICATED, "unauthenticated"),
+        RealRpcError(grpc.StatusCode.INTERNAL, "internal error"),
+    ],
+)
+def test_retry_logic_on_non_retryable_error(mock_client, non_retryable_error):
+    """Test that the query stops immediately on non-transient errors."""
+    query = TopicMessageQuery(topic_id="0.0.123").set_max_attempts(5).set_max_backoff(0.5)
+
+    mock_client.mirror_stub.subscribeTopic.side_effect = [non_retryable_error] * 5
+
+    on_error = MagicMock()
+    handle = query.subscribe(mock_client, on_message=MagicMock(), on_error=on_error)
+
+    handle._thread.join(timeout=1.0)
+
+    assert mock_client.mirror_stub.subscribeTopic.call_count == 1
+    on_error.assert_called_once_with(non_retryable_error)
+
+    assert not handle._thread.is_alive()
+
+
+def test_subscription_cancellation(mock_client):
+    """Test that cancelling a handle stops the subscription thread."""
+    query = TopicMessageQuery(topic_id="0.0.123")
+
+    def infinite_stream():
+        while True:
+            yield mirror_proto.ConsensusTopicResponse(message=b"ping")
+            time.sleep(0.1)
+
+    mock_call = MagicMock()
+    mock_call.__iter__.return_value = infinite_stream()
+
+    mock_client.mirror_stub.subscribeTopic.return_value = mock_call
+
+    on_message = MagicMock()
+    handle = query.subscribe(mock_client, on_message=on_message)
+
+    time.sleep(0.2)
+    assert handle._thread.is_alive()
+
+    handle.cancel()
+
+    handle._thread.join(timeout=1.0)
+
+    assert not handle._thread.is_alive()
+    mock_call.cancel.assert_called()