Allow custom error messages in snippet `waitForEvent`. (#989)
diff --git a/mobly/snippet/callback_handler_base.py b/mobly/snippet/callback_handler_base.py
index 95b5a4b..992d971 100644
--- a/mobly/snippet/callback_handler_base.py
+++ b/mobly/snippet/callback_handler_base.py
@@ -14,6 +14,7 @@
"""Module for the base class to handle Mobly Snippet Lib's callback events."""
import abc
import time
+from typing import Callable
from mobly.snippet import callback_event
from mobly.snippet import errors
@@ -178,7 +179,13 @@
raw_event = self.callEventWaitAndGetRpc(self._id, event_name, timeout)
return callback_event.from_dict(raw_event)
- def waitForEvent(self, event_name, predicate, timeout=None):
+ def waitForEvent(
+ self,
+ event_name: str,
+ predicate: Callable[[callback_event.CallbackEvent], bool],
+ timeout: float | None = None,
+ message: str | None = None,
+ ) -> callback_event.CallbackEvent:
"""Waits for an event of the specific name that satisfies the predicate.
This call will block until the expected event has been received or time
@@ -194,13 +201,13 @@
Args:
event_name: str, the name of the event to wait for.
- predicate: function, a function that takes an event (dictionary) and
- returns a bool.
+ predicate: function, the predicate used to test events.
timeout: float, the number of seconds to wait before giving up. If None,
it will be set to self.default_timeout_sec.
+ message: str, an optional error message to include if there is a timeout.
Returns:
- dictionary, the event that satisfies the predicate if received.
+ CallbackEvent, the event that satisfies the predicate if received.
Raises:
errors.CallbackHandlerTimeoutError: raised if no event that satisfies the
@@ -225,10 +232,11 @@
if predicate(event):
return event
+ custom_error = '' if message is None else f' Details: {message}.'
raise errors.CallbackHandlerTimeoutError(
self._device,
f'Timed out after {timeout}s waiting for an "{event_name}" event that '
- f'satisfies the predicate "{predicate.__name__}".',
+ f'satisfies the predicate "{predicate.__name__}".{custom_error}',
)
def getAll(self, event_name):
diff --git a/tests/mobly/snippet/callback_handler_base_test.py b/tests/mobly/snippet/callback_handler_base_test.py
index 32199b2..273df06 100644
--- a/tests/mobly/snippet/callback_handler_base_test.py
+++ b/tests/mobly/snippet/callback_handler_base_test.py
@@ -153,6 +153,24 @@
event = handler.waitForEvent('AsyncTaskResult', some_condition, 0.01)
self.assert_event_correct(event, MOCK_RAW_EVENT)
+ def test_wait_for_event_negative_with_custom_message(self):
+ handler = FakeCallbackHandler()
+ handler.mock_rpc_func.callEventWaitAndGetRpc = mock.Mock(
+ return_value=MOCK_RAW_EVENT
+ )
+
+ expected_msg = (
+ 'Timed out after 0.01s waiting for an "AsyncTaskResult" event that'
+ ' satisfies the predicate "<lambda>". Details: Failed for test.'
+ )
+
+ with self.assertRaisesRegex(
+ errors.CallbackHandlerTimeoutError, expected_msg
+ ):
+ handler.waitForEvent(
+ 'AsyncTaskResult', lambda _: False, 0.01, 'Failed for test'
+ )
+
def test_wait_for_event_negative(self):
handler = FakeCallbackHandler()
handler.mock_rpc_func.callEventWaitAndGetRpc = mock.Mock(
