ui_automator/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 functools
import logging
import os
import re
from typing import Any, Callable, Sequence

from absl import app
from absl import flags
import inflection
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

_GOOGLE_HOME_APP = {
    'id': 'gha',
    'name': 'Google Home Application (GHA)',
    'minVersion': '3.1.18.1',
    'minMatter': '231456000',
    'minThread': '231456000',
}
_MOBLY_SNIPPET_APK: str = 'com.chip.interop.moblysnippet'
_MOBLY_SNIPPET_APK_NAME: str = 'mbs'
_COMMISSIONING_FLAG_USAGE_GUIDE = (
    'Use --commission {DeviceName},{PairingCode},{GHARoom} to commission a'
    ' device to google fabric on GHA.'
)

_COMMISSION = flags.DEFINE_list(
    name='commission',
    default=None,
    help=_COMMISSIONING_FLAG_USAGE_GUIDE,
)


def _validate_commssioning_arg(
    device_name: str, pairing_code: str, gha_room: str
) -> None:
  """Returns None if commissioning values are valid.

  Args:
      device_name: Display name of commissioned device on GHA.
      pairing_code: An 11-digit or 21-digit numeric code which contains the
        information needed to commission a matter device.
      gha_room: Assigned room of commissioned device on GHA.

  Raises:
      ValueError: When any of passed arguments is invalid.
  """
  device_name_match = re.fullmatch(r'\S{1,24}', device_name)
  pairing_code_match = re.fullmatch(r'^(\d{11}|\d{21})$', pairing_code)
  gha_room_match = re.fullmatch(
      r'Attic|Back door|Backyard|Basement|Bathroom|Bedroom|Den|Dining'
      r' Room|Entryway|Family Room|Front door|Front'
      r' Yard|Garage|Hallway|Kitchen|Living Room|Master'
      r' Bedroom|Office|Shed|Side door',
      gha_room,
  )

  if not device_name_match:
    raise ValueError(
        'Value of DeviceName is invalid. Device name should be no more than 24'
        ' characters.'
    )
  if not pairing_code_match:
    raise ValueError(
        'Value of PairingCode is invalid. Paring code should be 11-digit or'
        ' 21-digit numeric code.'
    )
  if not gha_room_match:
    raise ValueError(
        'Value of GHARoom is invalid. Valid values for GHARoom: Attic|Back'
        ' door|Backyard|Basement|Bathroom|Bedroom|Den|Dining'
        ' Room|Entryway|Family Room|Front door|Front'
        ' Yard|Garage|Hallway|Kitchen|Living Room|Master'
        ' Bedroom|Office|Shed|Side door'
    )


flags.register_validator(
    'commission',
    lambda value: value and len(value) == 3,
    message=_COMMISSIONING_FLAG_USAGE_GUIDE,
)


def get_android_device_ready(func: Callable[..., Any]) -> Callable[..., Any]:
  """A decorator to check if an Android device can be controlled by installed apk.

  Args:
      func: The function to be wrapped.

  Returns:
      The wrapped function.

  Raises:
      NoAndroidDeviceError: When no device is connected.
      AndroidDeviceNotReadyError: When required snippet apk can not be loaded
      on device.
  """

  @functools.wraps(func)
  def wrapped_func(self, *args, **kwargs):
    try:
      self.load_device()
      self.load_snippet()
    except (
        errors.AndroidDeviceNotReadyError,
        errors.NoAndroidDeviceError,
    ) as e:
      raise errors.AndroidDeviceNotReadyError(f'{func.__name__} failed.') from e

    return func(self, *args, **kwargs)

  return wrapped_func


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))

  @get_android_device_ready
  def commission_device(
      self,
      device_name: str,
      pairing_code: str,
      gha_room: str
  ):
    """Commissions a device through installed apk `mbs` on Google Home App.

    Args:
        device_name: Display name of commissioned device on GHA.
        pairing_code: An 11-digit or 21-digit numeric code which contains the
          information needed to commission a matter device.
        gha_room: Assigned room of commissioned device on GHA.

    Raises:
        ValueError: When any of passed arguments is invalid.
        MoblySnippetError: When running `commissionDevice` method in snippet apk
        encountered an error.
    """
    self._logger.info('Start commissioning the device.')

    _validate_commssioning_arg(device_name, pairing_code, gha_room)

    device_name_snake_case = f'_{inflection.underscore(device_name)}'
    matter_device = {
        'id': device_name_snake_case,
        'name': device_name,
        'pairingCode': pairing_code,
        'roomName': gha_room,
    }

    try:
      self._connected_device.mbs.commissionDevice(
          _GOOGLE_HOME_APP, matter_device
      )
      self._logger.info('Successfully commission the device on GHA.')
    # TODO(b/298903492): Narrow exception type.
    except Exception as e:
      raise errors.MoblySnippetError(
          'Unable to continue automated commissioning process on'
          f' device({self._connected_device.device_info["serial"]}).'
      ) from 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',
    )


def _commission(ui_automator: UIAutomator) -> None:
  """Commissions a device to google fabric on GHA with flag values."""
  if _COMMISSION.value:
    ui_automator.commission_device(*_COMMISSION.value)


def _main(argv: Sequence[str]) -> None:
  del argv
  ui_automator = UIAutomator()
  _commission(ui_automator)


def run():
  app.run(_main)


if __name__ == '__main__':
  run()