venv/skylab_inventory/translation_utils.py

# -*- coding: utf-8 -*-
# Copyright 2018 The Chromium OS Authors. All rights reserved.
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.

"""Constants and util methods for the inventory."""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

from skylab_inventory.protos import common_pb2
from skylab_inventory.protos import device_pb2
from skylab_inventory.protos import server_pb2



def invert_map(origin_map):
  return {v: k for k, v in origin_map.iteritems()}


STATUS_MAP = {
    'primary': server_pb2.Server.Status.Value('STATUS_PRIMARY'),
    'backup': server_pb2.Server.Status.Value('STATUS_BACKUP'),
    'repair_required': server_pb2.Server.Status.Value('STATUS_REPAIR_REQUIRED'),
}

ROLE_MAP = {
    'afe': server_pb2.Server.Role.Value('ROLE_AFE'),
    'crash_server': server_pb2.Server.Role.Value('ROLE_CRASH_SERVER'),
    'database': server_pb2.Server.Role.Value('ROLE_DATABASE'),
    'database_slave': server_pb2.Server.Role.Value('ROLE_DATABASE_SLAVE'),
    'devserver': server_pb2.Server.Role.Value('ROLE_DEVSERVER'),
    'drone': server_pb2.Server.Role.Value('ROLE_DRONE'),
    'golo_proxy': server_pb2.Server.Role.Value('ROLE_GOLO_PROXY'),
    'host_scheduler': server_pb2.Server.Role.Value('ROLE_HOST_SCHEDULER'),
    'scheduler': server_pb2.Server.Role.Value('ROLE_SCHEDULER'),
    'sentinel': server_pb2.Server.Role.Value('ROLE_SENTINEL'),
    'shard': server_pb2.Server.Role.Value('ROLE_SHARD'),
    'suite_scheduler': server_pb2.Server.Role.Value('ROLE_SUITE_SCHEDULER'),
    'skylab_drone': server_pb2.Server.Role.Value('ROLE_SKYLAB_DRONE'),
    'skylab_suite_proxy': server_pb2.Server.Role.Value(
        'ROLE_SKYLAB_SUITE_PROXY'),
}

ENVIRONMENT_MAP = {
    'prod': common_pb2.Environment.Value('ENVIRONMENT_PROD'),
    'staging': common_pb2.Environment.Value('ENVIRONMENT_STAGING'),
}

SERVER_ATTRIBUTE_TYPE_MAP = {
    'ip': str,
    'max_processes': int,
    'mysql_server_id': str,
    'devserver_restricted_subnet': str,
    'devserver_port': int,
}

DUT_POOL_MAP = {
    'cq': device_pb2.SchedulableLabels.DUTPool.Value('DUT_POOL_CQ'),
    'bvt': device_pb2.SchedulableLabels.DUTPool.Value('DUT_POOL_BVT'),
    'suites': device_pb2.SchedulableLabels.DUTPool.Value('DUT_POOL_SUITES'),
    'cts': device_pb2.SchedulableLabels.DUTPool.Value('DUT_POOL_CTS'),
    'cts_prebuild': device_pb2.SchedulableLabels.DUTPool.Value(
        'DUT_POOL_CTS_PERBUILD'),
}
INVERTED_DUT_POOL_MAP = invert_map(DUT_POOL_MAP)

EC_TYPE_MAP = {
    'invalid': device_pb2.SchedulableLabels.ECType.Value('EC_TYPE_INVALID'),
    'chromeos': device_pb2.SchedulableLabels.ECType.Value('EC_TYPE_CHROME_OS'),
}
INVERTED_EC_TYPE_MAP = invert_map(EC_TYPE_MAP)

DEVICE_TYPE_TO_CLASS_MAP = {
    'duts': device_pb2.DeviceUnderTest,
    'servo_hosts': device_pb2.ServoHostDevice,
    'chamelons': device_pb2.ChameleonDevice,
}

CONNECTION_TYPES = {'servo_host_connections', 'chameleon_connections'}
DATA_SUBDIR = {'prod', 'staging', 'skylab'}
HARDWARE_CAPABILITIES_LABELS = {'modem', 'telephony', 'gpu_family', 'graphics',
                                'storage', 'power'}
TOPLEVEL_LABELS = {'platform', 'board', 'model', 'critical_pool',
                   'self_serve_pool', 'ec_type'}
SCHEDULABLE_LABELS = TOPLEVEL_LABELS | HARDWARE_CAPABILITIES_LABELS

# Each device should only have one label value for the unique label key.
UNIQUE_LABEL_KEY = ({'platform', 'board', 'model', 'ec_type'} |
                    HARDWARE_CAPABILITIES_LABELS)


class InvalidDataException(Exception):
  """Error raised when the data provided is invalid."""


class InvalidDataSubdir(InvalidDataException):
  """Error raised when the data_subdir is invalid."""


class InvalidEnvironment(InvalidDataException):
  """Error raised when environment is invalid."""


class InvalidEnvironmentSet(InvalidDataException):
  """Error raised when environment set is invalid."""


class InvalidRole(InvalidDataException):
  """Error raised when role is invalid."""


class InvalidStatus(InvalidDataException):
  """Error raised when status is invalid."""


class InvalidAttribute(InvalidDataException):
  """Error raised when attribute is invalid."""


class InvalidDeviceType(InvalidDataException):
  """Error raised when device_type is invalid."""


def validate_data_subdir(subdir):
  """Validate the subdir is one of DATA_SUBDIR.

  Args:
    subdir: The path to the data sub-directory.
  """
  if subdir not in DATA_SUBDIR:
    raise InvalidDataSubdir('Got %s. Data subdir must be one of %s. ' %
                            (subdir, DATA_SUBDIR))


def validate_environment(environment):
  """Validate that environment is one of ENVIRONMENT_MAP.keys.

  Args:
    environment: An environment of inventory data.

  Returns:
    Enum value for environment in ENVIRONMENT_MAP.
  """
  if environment not in ENVIRONMENT_MAP.keys():
    raise InvalidEnvironment('Got %s. environment must be one of %s. ' %
                             (environment, ENVIRONMENT_MAP.keys()))

  return ENVIRONMENT_MAP[environment]


def validate_environments(environments):
  """Validate environments is a subset of ENVIRONMENT_MAP.keys.

  Args:
    environments: A set of environments of inventory data.

  Returns:
    A set of enum values for |environments| in ENVIRONMENT_MAP.
  """
  if not isinstance(environments, set):
    raise InvalidEnvironmentSet('Got environments in %s type. environments must'
                                ' be a set.' % type(environments))

  environment_values = set()
  for e in environments:
    environment_values.add(validate_environment(e))

  return environment_values


def validate_server_role(role):
  """Validate server role is one of ROLE_MAP.keys.

  Args:
    role: server role.
  """
  if role not in ROLE_MAP.keys():
    raise InvalidRole('Got %s. role must be one of %s. ' %
                      (role, ROLE_MAP.keys()))


def validate_server_status(status):
  """Validate server status is one of STATUS_MAP.keys.

  Args:
    status: server status.
  """
  if status not in STATUS_MAP.keys():
    raise InvalidStatus('Got %s. status must be one of %s.' %
                        (status, STATUS_MAP.keys()))


def validate_server_attribute(attribute):
  """Validate server attribute is one of SERVER_ATTRIBUTE_TYPE_MAP.keys.

  Args:
    attribute: The key of server attribute.
  """
  if attribute not in SERVER_ATTRIBUTE_TYPE_MAP:
    raise InvalidAttribute('Got %s. attribute key must be one of %s. ' %
                           (attribute, SERVER_ATTRIBUTE_TYPE_MAP.keys()))


def validate_device_type(device_type):
  """Validate device_type is one of DEVICE_TYPE_TO_CLASS_MAP.keys.

  Args:
    device_type: a type of device.
  """
  if device_type not in DEVICE_TYPE_TO_CLASS_MAP:
    raise InvalidDeviceType('Got %s. device_type must be one of %s' %
                            (device_type, DEVICE_TYPE_TO_CLASS_MAP.keys()))