ui_automator.py

"""UI Automator controller in python layer.

A python controller that can trigger mobly UI automator snippet to achieve some
automated UI operations on Android phones.
"""
import logging
import os

from mobly import utils
from mobly.controllers import android_device
from mobly.controllers.android_device_lib import adb
from mobly.snippet import errors as snippet_errors

from ui_automator import errors

_MOBLY_SNIPPET_APK: str = 'com.chip.interop.moblysnippet'
_MOBLY_SNIPPET_APK_NAME: str = 'mbs'


class UIAutomator:
  """UI Automator controller in python layer."""

  def __init__(self, logger: logging.Logger = logging.getLogger()):
    """Inits UIAutomator.

    Interacts with android device by pre-installed apk. Can perform
    methods provided by a snippet installed on the android device.

    Args:
        logger: Injected logger, if None, specify the default(root) one.
    """
    self._logger = logger
    self._connected_device: android_device.AndroidDevice | None = None

  def load_device(self):
    """Selects the first connected android device.

    Raises:
        NoAndroidDeviceError: When no device is connected.
    """
    try:
      android_devices = android_device.get_all_instances()
    except adb.AdbError as exc:
      raise errors.AdbError(
          'Please install adb and add it to your PATH environment variable.'
      ) from exc

    if not android_devices:
      raise errors.NoAndroidDeviceError(
          'No Android device connected to the host computer.'
      )
    self._connected_device = android_devices[0]
    self._logger.info(f'connected device: [{self._connected_device}]')

  def load_snippet(self):
    """Loads needed mobly snippet installed on android device.

    Raises:
        NoAndroidDeviceError: When no device is connected.
        AndroidDeviceNotReadyError: When required snippet apk can not be loaded
        on device.
    """
    if not self._connected_device:
      raise errors.NoAndroidDeviceError(
          'No Android device connected to the host computer.'
      )

    if not self._is_apk_installed(self._connected_device, _MOBLY_SNIPPET_APK):
      self._install_apk(self._connected_device, self._get_mbs_apk_path())

    try:
      self._connected_device.load_snippet(
          _MOBLY_SNIPPET_APK_NAME, _MOBLY_SNIPPET_APK
      )
    except snippet_errors.ServerStartPreCheckError as exc:
      # Raises when apk not be installed or instrumented on device.
      raise errors.AndroidDeviceNotReadyError(
          f"Check device({self._connected_device.device_info['serial']}) has"
          ' installed required apk.'
      ) from exc
    except (
        snippet_errors.ServerStartError,
        snippet_errors.ProtocolError,
    ) as exc:
      raise errors.AndroidDeviceNotReadyError(
          f"Check device({self._connected_device.device_info['serial']}) can"
          ' load required apk.'
      ) from exc
    except android_device.SnippetError as e:
      # Error raises when registering a package twice or a package with
      # duplicated name. Thus, It is okay to pass here as long as the snippet
      # has been loaded.
      self._logger.debug(str(e))

  def _is_apk_installed(
      self, device: android_device.AndroidDevice, package_name: str
  ) -> bool:
    """Checks if given package is already installed on the Android device.

    Args:
      device: An AndroidDevice object.
      package_name: The Android app package name.

    Raises:
      adb.AdbError: When executing adb command fails.

    Returns:
      True if package is installed. False otherwise.
    """
    out = device.adb.shell(['pm', 'list', 'package'])
    return bool(utils.grep('^package:%s$' % package_name, out))

  def _install_apk(
      self, device: android_device.AndroidDevice, apk_path: str
  ) -> None:
    """Installs required apk snippet to the given Android device.

    Args:
      device: An AndroidDevice object.
      apk_path: The absolute file path where the apk is located.
    """
    device.adb.install(['-r', '-g', apk_path])

  def _get_mbs_apk_path(self) -> str:
    return os.path.join(
        os.path.dirname(os.path.abspath(__file__)),
        'android',
        'app',
        'snippet-0.0.0.apk',
    )


if __name__ == '__main__':
  UIAutomator().load_device()