text: # Copyright 2022 Google Inc.
  #
  # Licensed under the Apache License, Version 2.0 (the "License");
  # you may not use this file except in compliance with the License.
  # You may obtain a copy of the License at
  #
  #     http://www.apache.org/licenses/LICENSE-2.0
  #
  # Unless required by applicable law or agreed to in writing, software
  # distributed under the License is distributed on an "AS IS" BASIS,
  # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  # See the License for the specific language governing permissions and
  # limitations under the License.
  """The callback handler V2 module for Android Mobly Snippet Lib."""
  
  from mobly.snippet import callback_handler_base
  from mobly.snippet import errors
  
  # The timeout error meesage when pulling events from the server
  TIMEOUT_ERROR_MESSAGE = 'EventSnippetException: timeout.'
  
  
  class CallbackHandlerV2(callback_handler_base.CallbackHandlerBase):
    """The callback handler V2 class for Android Mobly Snippet Lib."""
  
    def callEventWaitAndGetRpc(self, callback_id, event_name, timeout_sec):
      """Waits and returns an existing CallbackEvent for the specified identifier.
  
      This function calls snippet lib's eventWaitAndGet RPC.
  
      Args:
        callback_id: str, the callback identifier.
        event_name: str, the callback name.
        timeout_sec: float, the number of seconds to wait for the event.
  
      Returns:
        The event dictionary.
  
      Raises:
        errors.CallbackHandlerTimeoutError: The expected event does not occur
          within the time limit.
      """
      timeout_ms = int(timeout_sec * 1000)
      try:
        return self._event_client.eventWaitAndGet(callback_id, event_name,
                                                  timeout_ms)
      except Exception as e:
        if TIMEOUT_ERROR_MESSAGE in str(e):
          raise errors.CallbackHandlerTimeoutError(
              self._device, (f'Timed out after waiting {timeout_sec}s for event '
                             f'"{event_name}" triggered by {self._method_name} '
                             f'({self.callback_id}).')) from e
        raise
  
    def callEventGetAllRpc(self, callback_id, event_name):
      """Gets all existing events for the specified identifier without waiting.
  
      This function calls snippet lib's eventGetAll RPC.
  
      Args:
        callback_id: str, the callback identifier.
        event_name: str, the callback name.
  
      Returns:
        A list of event dictionaries.
      """
      return self._event_client.eventGetAll(callback_id, event_name)
  