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 JSON RPC client base for communicating with snippet servers.
  
  The JSON RPC protocol expected by this module is:
  
  .. code-block:: json
  
    Request:
    {
      'id': <Required. Monotonically increasing integer containing the ID of this
            request.>,
      'method': <Required. String containing the name of the method to execute.>,
      'params': <Required. JSON array containing the arguments to the method,
                `null` if no positional arguments for the RPC method.>,
      'kwargs': <Optional. JSON dict containing the keyword arguments for the
                method, `null` if no positional arguments for the RPC method.>,
    }
  
    Response:
    {
      'error': <Required. String containing the error thrown by executing the
               method, `null` if no error occurred.>,
      'id': <Required. Int id of request that this response maps to.>,
      'result': <Required. Arbitrary JSON object containing the result of
                executing the method, `null` if the method could not be executed
                or returned void.>,
      'callback': <Required. String that represents a callback ID used to
                  identify events associated with a particular CallbackHandler
                  object, `null` if this is not an asynchronous RPC.>,
    }
  """
  
  import abc
  import json
  import threading
  import time
  
  from mobly.snippet import errors
  
  # Maximum logging length of RPC response in DEBUG level when verbose logging is
  # off.
  _MAX_RPC_RESP_LOGGING_LENGTH = 1024
  
  # The required field names of RPC response.
  RPC_RESPONSE_REQUIRED_FIELDS = ('id', 'error', 'result', 'callback')
  
  
  class ClientBase(abc.ABC):
    """Base class for JSON RPC clients that connect to snippet servers.
  
    Connects to a remote device running a JSON RPC compatible server. Users call
    the function `start_server` to start the server on the remote device before
    sending any RPC. After sending all RPCs, users call the function `stop`
    to stop the snippet server and release all the requested resources.
  
    Attributes:
      package: str, the user-visible name of the snippet library being
        communicated with.
      log: Logger, the logger of the corresponding device controller.
      verbose_logging: bool, if True, prints more detailed log
        information. Default is True.
    """
  
    def __init__(self, package, device):
      """Initializes the instance of ClientBase.
  
      Args:
        package: str, the user-visible name of the snippet library being
          communicated with.
        device: DeviceController, the device object associated with a client.
      """
  
      self.package = package
      self.log = device.log
      self.verbose_logging = True
      self._device = device
      self._counter = None
      self._lock = threading.Lock()
      self._event_client = None
  
    def __del__(self):
      self.close_connection()
  
    def initialize(self):
      """Initializes the snippet client to interact with the remote device.
  
      This function contains following stages:
        1. preparing to start the snippet server.
        2. starting the snippet server on the remote device.
        3. making a connection to the snippet server.
  
      If error occurs at any stage, this function will abort the initialization
      process and call `stop` to clean up.
  
      Raises:
        errors.ProtocolError: something went wrong when exchanging data with the
          server.
        errors.ServerStartPreCheckError: when prechecks for starting the server
          failed.
        errors.ServerStartError: when failed to start the snippet server.
      """
  
      # Use log.info here so people can follow along with the initialization
      # process. Initialization can be slow, especially if there are
      # multiple snippets, this avoids the perception that the framework
      # is hanging for a long time doing nothing.
      self.log.info('Initializing the snippet package %s.', self.package)
      start_time = time.perf_counter()
  
      try:
        self.log.debug('Preparing to start the snippet server of %s.',
                       self.package)
        self.before_starting_server()
  
        self.log.debug('Starting the snippet server of %s.', self.package)
        self.start_server()
  
        self.log.debug('Making a connection to the snippet server of %s.',
                       self.package)
        self._make_connection()
  
      except Exception:
        self.log.error(
            'Error occurred trying to start and connect to the snippet server '
            'of %s.', self.package)
        try:
          self.stop()
        except Exception:  # pylint: disable=broad-except
          # Only prints this exception and re-raises the original exception
          self.log.exception(
              'Failed to stop the snippet package %s after failure to start '
              'and connect.', self.package)
  
        raise
  
      self.log.debug('Snippet package %s initialized after %.1fs.', self.package,
                     time.perf_counter() - start_time)
  
    @abc.abstractmethod
    def before_starting_server(self):
      """Performs the preparation steps before starting the remote server.
  
      For example, subclass can check or modify the device settings at this
      stage.
  
      Raises:
        errors.ServerStartPreCheckError: when prechecks for starting the server
          failed.
      """
  
    @abc.abstractmethod
    def start_server(self):
      """Starts the server on the remote device.
  
      The client has completed the preparations, so the client calls this
      function to start the server.
      """
  
    def _make_connection(self):
      """Proxy function of make_connection.
  
      This function resets the RPC id counter before calling `make_connection`.
      """
      self._counter = self._id_counter()
      self.make_connection()
  
    @abc.abstractmethod
    def make_connection(self):
      """Makes a connection to the snippet server on the remote device.
  
      This function makes a connection to the server and sends a handshake
      request to ensure the server is available for upcoming RPCs.
  
      There are two types of connections used by snippet clients:
      * The client makes a new connection each time it needs to send an RPC.
      * The client makes a connection in this stage and uses it for all the RPCs.
        In this case, the client should implement `close_connection` to close
        the connection.
  
      Raises:
        errors.ProtocolError: something went wrong when exchanging data with the
          server.
      """
  
    def __getattr__(self, name):
      """Wrapper for python magic to turn method calls into RPCs."""
  
      def rpc_call(*args, **kwargs):
        return self._rpc(name, *args, **kwargs)
  
      return rpc_call
  
    def _id_counter(self):
      """Returns an id generator."""
      i = 0
      while True:
        yield i
        i += 1
  
    def set_snippet_client_verbose_logging(self, verbose):
      """Switches verbose logging. True for logging full RPC responses.
  
      By default it will write full messages returned from RPCs. Turning off the
      verbose logging will result in writing no more than
      _MAX_RPC_RESP_LOGGING_LENGTH characters per RPC returned string.
  
      _MAX_RPC_RESP_LOGGING_LENGTH will be set to 1024 by default. The length
      contains the full RPC response in JSON format, not just the RPC result
      field.
  
      Args:
        verbose: bool, if True, turns on verbose logging, otherwise turns off.
      """
      self.log.info('Sets verbose logging to %s.', verbose)
      self.verbose_logging = verbose
  
    @abc.abstractmethod
    def restore_server_connection(self, port=None):
      """Reconnects to the server after the device was disconnected.
  
      Instead of creating a new instance of the client:
        - Uses the given port (or finds a new available host port if 0 or None is
        given).
        - Tries to connect to the remote server with the selected port.
  
      Args:
        port: int, if given, this is the host port from which to connect to the
          remote device port. Otherwise, finds a new available port as host
          port.
  
      Raises:
        errors.ServerRestoreConnectionError: when failed to restore the connection
          to the snippet server.
      """
  
    def _rpc(self, rpc_func_name, *args, **kwargs):
      """Sends an RPC to the server.
  
      Args:
        rpc_func_name: str, the name of the snippet function to execute on the
          server.
        *args: any, the positional arguments of the RPC request.
        **kwargs: any, the keyword arguments of the RPC request.
  
      Returns:
        The result of the RPC.
  
      Raises:
        errors.ProtocolError: something went wrong when exchanging data with the
          server.
        errors.ApiError: the RPC went through, however executed with errors.
      """
      try:
        self.check_server_proc_running()
      except Exception:
        self.log.error(
            'Server process running check failed, skip sending RPC method(%s).',
            rpc_func_name)
        raise
  
      with self._lock:
        rpc_id = next(self._counter)
        request = self._gen_rpc_request(rpc_id, rpc_func_name, *args, **kwargs)
  
        self.log.debug('Sending RPC request %s.', request)
        response = self.send_rpc_request(request)
        self.log.debug('RPC request sent.')
  
        if self.verbose_logging or _MAX_RPC_RESP_LOGGING_LENGTH >= len(response):
          self.log.debug('Snippet received: %s', response)
        else:
          self.log.debug('Snippet received: %s... %d chars are truncated',
                         response[:_MAX_RPC_RESP_LOGGING_LENGTH],
                         len(response) - _MAX_RPC_RESP_LOGGING_LENGTH)
  
      response_decoded = self._decode_response_string_and_validate_format(
          rpc_id, response)
      return self._handle_rpc_response(rpc_func_name, response_decoded)
  
    @abc.abstractmethod
    def check_server_proc_running(self):
      """Checks whether the server is still running.
  
      If the server is not running, it throws an error. As this function is called
      each time the client tries to send an RPC, this should be a quick check
      without affecting performance. Otherwise it is fine to not check anything.
  
      Raises:
        errors.ServerDiedError: if the server died.
      """
  
    def _gen_rpc_request(self, rpc_id, rpc_func_name, *args, **kwargs):
      """Generates the JSON RPC request.
  
      In the generated JSON string, the fields are sorted by keys in ascending
      order.
  
      Args:
        rpc_id: int, the id of this RPC.
        rpc_func_name: str, the name of the snippet function to execute
          on the server.
        *args: any, the positional arguments of the RPC.
        **kwargs: any, the keyword arguments of the RPC.
  
      Returns:
        A string of the JSON RPC request.
      """
      data = {'id': rpc_id, 'method': rpc_func_name, 'params': args}
      if kwargs:
        data['kwargs'] = kwargs
      return json.dumps(data, sort_keys=True)
  
    @abc.abstractmethod
    def send_rpc_request(self, request):
      """Sends the JSON RPC request to the server and gets a response.
  
      Note that the request and response are both in string format. So if the
      connection with server provides interfaces in bytes format, please
      transform them to string in the implementation of this function.
  
      Args:
        request: str, a string of the RPC request.
  
      Returns:
        A string of the RPC response.
  
      Raises:
        errors.ProtocolError: something went wrong when exchanging data with the
          server.
      """
  
    def _decode_response_string_and_validate_format(self, rpc_id, response):
      """Decodes response JSON string to python dict and validates its format.
  
      Args:
        rpc_id: int, the actual id of this RPC. It should be the same with the id
          in the response, otherwise throws an error.
        response: str, the JSON string of the RPC response.
  
      Returns:
        A dict decoded from the response JSON string.
  
      Raises:
        errors.ProtocolError: if the response format is invalid.
      """
      if not response:
        raise errors.ProtocolError(self._device,
                                   errors.ProtocolError.NO_RESPONSE_FROM_SERVER)
  
      result = json.loads(response)
      for field_name in RPC_RESPONSE_REQUIRED_FIELDS:
        if field_name not in result:
          raise errors.ProtocolError(
              self._device,
              errors.ProtocolError.RESPONSE_MISSING_FIELD % field_name)
  
      if result['id'] != rpc_id:
        raise errors.ProtocolError(self._device,
                                   errors.ProtocolError.MISMATCHED_API_ID)
  
      return result
  
    def _handle_rpc_response(self, rpc_func_name, response):
      """Handles the content of RPC response.
  
      If the RPC response contains error information, it throws an error. If the
      RPC is asynchronous, it creates and returns a callback handler
      object. Otherwise, it returns the result field of the response.
  
      Args:
        rpc_func_name: str, the name of the snippet function that this RPC
          triggered on the snippet server.
        response: dict, the object decoded from the response JSON string.
  
      Returns:
        The result of the RPC. If synchronous RPC, it is the result field of the
        response. If asynchronous RPC, it is the callback handler object.
  
      Raises:
        errors.ApiError: if the snippet function executed with errors.
      """
  
      if response['error']:
        raise errors.ApiError(self._device, response['error'])
      if response['callback'] is not None:
        return self.handle_callback(response['callback'], response['result'],
                                    rpc_func_name)
      return response['result']
  
    @abc.abstractmethod
    def handle_callback(self, callback_id, ret_value, rpc_func_name):
      """Creates a callback handler for the asynchronous RPC.
  
      Args:
        callback_id: str, the callback ID for creating a callback handler object.
        ret_value: any, the result field of the RPC response.
        rpc_func_name: str, the name of the snippet function executed on the
          server.
  
      Returns:
        The callback handler object.
      """
  
    @abc.abstractmethod
    def stop(self):
      """Releases all the resources acquired in `initialize`."""
  
    @abc.abstractmethod
    def close_connection(self):
      """Closes the connection to the snippet server on the device.
  
      This is a unilateral closing from the client side, without tearing down
      the snippet server running on the device.
  
      The connection to the snippet server can be re-established by calling
      `restore_server_connection`.
      """
  