catatester/devicemap/devicemap_lib.py

"""Utility methods to retrieve devicemap.yaml data."""

import json
import logging
import os
import yaml


class Error(Exception):
  pass


class DeviceMapError(Error):
  pass


class Configs(object):
  """Represents all of the configs in the devicemap."""
  _DEVICE_NODE_NAME = 'hosted_devices'
  _HOST_IP = 'host_ip'

  def __init__(self, configs_dict=None):
    self._configs_dict = configs_dict or _ReadDeviceMap()
    self.allowed_dimensions = self._configs_dict.get('allowed_dimensions')
    self.hosts = self._configs_dict.get('hosts')

  def GetAllowedDimensions(self, config_type):
    """Gets the set() of allowed dimensions for the given type of config."""
    if config_type not in self.allowed_dimensions:
      raise DeviceMapError('config type not defined: %s' % config_type)
    return set(self.allowed_dimensions[config_type])

  def GetBotConfig(self, hostname, bot_num):
    """Gets the bot config for the given hostname & bot_num (1-indexed)."""
    host = self._GetHostConfig(hostname)
    hosted_devices = host[self._DEVICE_NODE_NAME]
    device_index = bot_num - 1
    if device_index >= len(hosted_devices):
      raise DeviceMapError('Bot #%d is not defined under host %s' % (
        bot_num, hostname))
    config_dict = hosted_devices[device_index]
    host_dimensions = self._GetHostDimensions(host)
    return BotConfig(config_dict, host_dimensions)

  def GetNumBotConfigs(self, hostname):
    """Gets the # of bot configs assigned to the given hostname."""
    return len(self._GetHostedDevices(hostname))

  def _GetHostConfig(self, hostname):
    """Returns the complete host config dict for a given hostname."""
    if hostname not in self.hosts:
      raise DeviceMapError('host %s is not defined in devicemap' % hostname)
    return self.hosts[hostname]

  def _GetHostDimensions(self, host_config):
    """Returns a dict of bot dimensions, populated from the host.

    These dimensions can be overridden at the individual device level.
    """
    return {
      'location': host_config.get('location', 'UNSPECIFIED')
    }

  def _GetHostedDevices(self, hostname):
    return self._GetHostConfig(hostname)[self._DEVICE_NODE_NAME]

  def GetHosts(self):
    return self.hosts

  def GetHostIps(self):
    # Implicitly making host_ip a mandatory field here
    return [self.hosts[h][self._HOST_IP] for h in self.hosts.keys()]


class BotConfig(object):
  """Represents a single bot config for a Testbed or Device."""

  def __init__(self, bot_config_dict, host_dimensions=None):
    """Initialize bot config object.

    Args:
      bot_config_dict: Dict of bot info for a single device from devicemap.
      host_dimensions: Bot dimensions populated from the host of the device.
    """
    self.type = self._ParseConfigType(bot_config_dict)
    self._raw_type = list(bot_config_dict.keys())[0]
    self._data = bot_config_dict[self._raw_type]
    self._host_dimensions = host_dimensions or {}
    if self.type == 'Testbed':
      self.dimensions, self.state = self._ParseTestbedConfig()
    elif self.type == 'Device':
      self.dimensions, self.state = self._ParseDeviceConfig()
    else:
      raise NotImplementedError('devicemap type "%s" not defined' % self.type)

  def __repr__(self):
    return str(self.as_dict)

  @property
  def as_dict(self):
    # Transform dimensions & state to conform to swarming's schema for bots
    for dimension, value in self.dimensions.items():
      if not value:
        raise DeviceMapError('Dimensions must not be empty: %s' % dimension)
      # Dimension values must be a list/tuple
      if not hasattr(value, '__iter__'):
        self.dimensions[dimension] = [value]
    return {
      'dimensions': self.dimensions,
      'state': self.state,
    }

  def _ParseConfigType(self, config_dict):
    if not len(list(config_dict.keys())) == 1:
      raise DeviceMapError(
        ('Configuration invalid -- expected single entry for a "Testbed" '
         'or a single device. Got: %s') % list(config_dict.keys()))
    config_type = list(config_dict.keys())[0]
    # Any configs that are named all in lowercase are assumed to be a device.
    # (e.g. lexx, luther, biggie, joplin, etc). Any configs with a leading
    # uppercase are a specified config type (e.g. Testbed).
    if config_type.islower():
      config_type = 'Device'
    return config_type

  def _ParseTestbedConfig(self):
    """Get the bot config dict() for the given hostname & bot_num."""
    dimensions = self._host_dimensions
    dimensions.update(self._data.get('dimensions'))
    state = self._data.get('extra_state')
    return dimensions, state

  def _ParseDeviceConfig(self):
    """Get the bot config dict() for the given hostname & bot_num."""
    state = {}
    dimensions = self._host_dimensions
    dimensions.update(self._data)
    dimensions['device_type'] = self._raw_type
    dimensions['group'] = dimensions.get('group') or 'COMMON'
    dimensions['nic'] = dimensions.get('nic') or 'ethernet'
    dimensions['use_tunnel'] = dimensions.get('use_tunnel') or 'false'
    # FDR device between tests unless in 'signed_in' group.
    dimensions['factory_reset'] = (
        dimensions.get('factory_reset') or
        str('signed_in' not in dimensions['group']).lower())
    # Rename 'ip' to 'device_ip' for convenience (keeps config file concise)
    if 'ip' in dimensions:
      dimensions['device_ip'] = dimensions['ip']
      del dimensions['ip']
    if 'extra_state' in dimensions:
      state = dimensions.get('extra_state', {})
      del dimensions['extra_state']
    return dimensions, state


def _ReadDeviceMap():
  """Reads all configs from devicemap.yaml config file as a dict()."""
  pwd = os.path.dirname(os.path.realpath(__file__))
  with open(os.path.join(pwd, 'devicemap.yaml'), 'r') as f:
    return yaml.safe_load(f)