src/hiero_sdk_python/query/topic_message_query.py (1)

61-62: ⚠️ Potential issue | 🟡 Minor | ⚡ Quick win

Type annotation mismatch: _error_handler declared as Callable[[], None] but used with Exception argument.

Line 62 declares _error_handler as Callable[[], None], but:

• set_error_handler (line 85) expects Callable[[Exception], None]
• _on_error (line 118) takes err: Exception
• The handler is invoked with self._error_handler(e) at line 243

This inconsistency means type checkers will report errors, and users following the type hint at line 62 would provide incompatible handlers.

🐛 Proposed fix

         self._completion_handler: Callable[[], None] | None = self._on_complete
-        self._error_handler: Callable[[], None] | None = self._on_error
+        self._error_handler: Callable[[Exception], None] | None = self._on_error

tests/unit/topic_message_query_test.py (3)

152-154: ⚠️ Potential issue | 🟡 Minor | ⚡ Quick win

Chunk ordering is not verified — test would pass even if chunks are assembled in wrong order.

Using in only checks presence, not position. If the implementation incorrectly orders chunks, this test won't detect it.

🐛 Proposed fix to verify correct ordering

     assert len(received_messages) == 1
-    assert b"chunk-1" in received_messages[0].contents
-    assert b"chunk-2" in received_messages[0].contents
+    # Verify chunks are assembled in correct order
+    contents = received_messages[0].contents
+    chunk1_pos = contents.find(b"chunk-1")
+    chunk2_pos = contents.find(b"chunk-2")
+    assert chunk1_pos != -1, "chunk-1 not found in assembled message"
+    assert chunk2_pos != -1, "chunk-2 not found in assembled message"
+    assert chunk1_pos < chunk2_pos, "Chunks assembled in wrong order: chunk-1 should come before chunk-2"

As per coding guidelines, tests should catch regressions in chunk ordering behavior.

---

167-178: ⚠️ Potential issue | 🟡 Minor | ⚡ Quick win

Missing assertion that on_error is invoked after retries exhaust.

When max_attempts are exhausted with retryable errors, on_error should be called. The test creates an on_error=MagicMock() but never verifies it was invoked, so regressions in terminal error handling would go undetected.

🐛 Proposed fix

 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())
+    on_error = MagicMock()
+    handle = query.subscribe(mock_client, on_message=MagicMock(), on_error=on_error)

     handle._thread.join(timeout=2.0)
+    assert not handle._thread.is_alive(), "Thread should have terminated after retries exhausted"

     assert mock_client.mirror_stub.subscribeTopic.call_count == 2
+    on_error.assert_called_once_with(error)

As per coding guidelines, "Tests must provide useful error messages when they fail for future debugging."

---

220-224: ⚠️ Potential issue | 🟡 Minor | ⚡ Quick win

Timing-dependent sleep can cause flaky tests in CI.

time.sleep(0.2) followed by assert handle._thread.is_alive() is scheduler-dependent. In slow CI environments, the thread may not have started processing yet, or conversely, the sleep may be insufficient. Use an event to synchronize on first message receipt before cancellation.

🐛 Proposed fix using event synchronization

+import threading
+
 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()
+    first_message_seen = threading.Event()
+    on_message = MagicMock(side_effect=lambda _: first_message_seen.set())
     handle = query.subscribe(mock_client, on_message=on_message)

-    time.sleep(0.2)
-    assert handle._thread.is_alive()
+    assert first_message_seen.wait(timeout=2.0), "Expected at least one streamed message before cancellation"
+    assert handle._thread.is_alive(), "Subscription thread should be alive before cancel"

     handle.cancel()

As per coding guidelines, "No timing-dependent or unseeded random assertions."

tests/unit/topic_message_query_test.py (4)

383-408: ⚠️ Potential issue | 🟠 Major | ⚡ Quick win

Replace timing-dependent sleep with event-based synchronization.

The time.sleep(0.2) at line 400 followed by is_alive() assertion at line 401 is scheduler-dependent and can flap in CI under high load or slow schedulers. Use an event to deterministically wait for the first message delivery before asserting thread liveness and performing cancellation.

⏱️ Proposed fix using threading.Event

+import threading
 import time
 from datetime import datetime, timezone
@@
 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()
+    first_message_seen = threading.Event()
+    on_message = MagicMock(side_effect=lambda _: first_message_seen.set())
     handle = query.subscribe(mock_client, on_message=on_message)

-    time.sleep(0.2)
-    assert handle._thread.is_alive()
+    assert first_message_seen.wait(timeout=1.0), "Expected at least one message before cancellation"
+    assert handle._thread.is_alive(), "Subscription thread should be alive before cancel()"

     handle.cancel()

     handle._thread.join(timeout=1.0)

     assert not handle._thread.is_alive()
     mock_call.cancel.assert_called()

As per coding guidelines, "No timing-dependent or unseeded random assertions."

---

54-80: ⚠️ Potential issue | 🟡 Minor | ⚡ Quick win

Verify fluent interface contract: setters must return self.

The test validates internal state but doesn't assert that each setter returns self, which is required for method chaining. If a setter is accidentally changed to return None, chaining would break but this test would pass.

🔗 Proposed fix

 def test_topic_message_query_initialization():
     """Test initializing the query with various parameter types and setters."""
     start = datetime(2023, 1, 1, tzinfo=timezone.utc)

     def mock_complete():
         pass

     def mock_error(e):
         pass

-    query = (
-        TopicMessageQuery()
-        .set_topic_id("0.0.123")
-        .set_start_time(start)
-        .set_limit(5)
-        .set_chunking_enabled(True)
-        .set_completion_handler(mock_complete)
-        .set_error_handler(mock_error)
-    )
+    query = TopicMessageQuery()
+    assert query.set_topic_id("0.0.123") is query, "set_topic_id should return self"
+    assert query.set_start_time(start) is query, "set_start_time should return self"
+    assert query.set_limit(5) is query, "set_limit should return self"
+    assert query.set_chunking_enabled(True) is query, "set_chunking_enabled should return self"
+    assert query.set_completion_handler(mock_complete) is query, "set_completion_handler should return self"
+    assert query.set_error_handler(mock_error) is query, "set_error_handler should return self"

     assert query._topic_id.topicNum == 123
     assert query._start_time.seconds == int(start.timestamp())
     assert query._limit == 5
     assert query._chunking_enabled is True
     assert query._completion_handler == mock_complete
     assert query._error_handler == mock_error

As per coding guidelines, "Assert fluent setters return self".

---

304-306: ⚠️ Potential issue | 🟡 Minor | ⚡ Quick win

Chunk ordering is not verified — test passes even if chunks are assembled backwards.

The test uses in at lines 305-306, which only checks presence. If the implementation concatenates chunk-2 before chunk-1, this assertion would still pass.

🔍 Proposed fix to verify correct ordering

     assert len(received_messages) == 1
-    assert b"chunk-1" in received_messages[0].contents
-    assert b"chunk-2" in received_messages[0].contents
+    # Verify chunks are assembled in correct order
+    contents = received_messages[0].contents
+    chunk1_pos = contents.find(b"chunk-1")
+    chunk2_pos = contents.find(b"chunk-2")
+    assert chunk1_pos != -1, "chunk-1 not found in assembled message"
+    assert chunk2_pos != -1, "chunk-2 not found in assembled message"
+    assert chunk1_pos < chunk2_pos, "Chunks assembled in wrong order: chunk-1 must precede chunk-2"

As per coding guidelines, tests must catch regressions in ordering behavior.

---

345-355: ⚠️ Potential issue | 🟡 Minor | ⚡ Quick win

Verify terminal error callback and thread termination after retries exhaust.

The test passes on_error=MagicMock() at line 350 but never asserts it's invoked when max_attempts are exhausted. The implementation should call on_error with the final error before stopping. Additionally, join(timeout=2.0) at line 352 should be followed by an assertion that the thread actually terminated—if it times out, the subscription count assertion at line 354 may pass spuriously.

🧪 Proposed fix

 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())
+    on_error = MagicMock()
+    handle = query.subscribe(mock_client, on_message=MagicMock(), on_error=on_error)

     handle._thread.join(timeout=2.0)
+    assert not handle._thread.is_alive(), "Subscription thread should have terminated after retries exhausted"

     assert mock_client.mirror_stub.subscribeTopic.call_count == 2
+    on_error.assert_called_once_with(error)

As per coding guidelines, "Tests must provide useful error messages when they fail for future debugging" and should verify error handler contracts.

src/hiero_sdk_python/query/topic_message_query.py (1)

176-177: ⚠️ Potential issue | 🟠 Major

Don't let a resumed query fall back to unlimited streaming.

Proto field: ConsensusTopicQuery.limit (#4). Issue type: Wrong default. Once state.count reaches the caller's cap, the retry path sends limit = 0, but the mirror schema defines zero/unset as "return messages indefinitely", so a reconnect after the cap is reached can over-deliver instead of stopping. Stop locally once state.count >= self._limit, or avoid emitting a retry request with limit=0. (github.com)
As per coding guidelines, "Compare the SDK class against the proto schema."

Also applies to: 227-259

src/hiero_sdk_python/utils/subscription_handle.py (1)

20-26: ⚠️ Potential issue | 🟠 Major

Cancel the active call after releasing _lock.

Both _set_call() and cancel() invoke call.cancel() inside the critical section. That makes the lock cover an external call and can stall the other thread trying to register or cancel the stream. Capture the call reference while locked, then invoke cancel() once the mutex is released.

Also applies to: 28-34

tests/unit/topic_message_query_test.py (1)

397-405: ⚠️ Potential issue | 🟡 Minor

Make this cancellation test event-driven.

time.sleep(0.2) plus an immediate liveness check is scheduler-dependent and can flap in CI. Wait for the first streamed message or for subscribeTopic() to be entered before asserting the thread is alive and calling cancel(). As per coding guidelines, "No timing-dependent or unseeded random assertions."