venv/skylab_inventory/scripts/print_dut_dimensions.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.

"""Script to print swarming dimensions for a given DUT."""

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

from skylab_inventory import text_manager
from skylab_inventory.protos import common_pb2
from skylab_inventory.translation import swarming

import argparse
import json


def _dimensions_json(dut):
  """Returns the dimensions for swarming bot as json."""
  common = dut.common
  dims = {}
  # These are required.
  # We must rename this dimension to dut_id, since 'id' is a magic dimension
  # that must equal the swraming bot id. See
  # https://chromium.googlesource.com/infra/luci/luci-py/+/master/appengine/swarming/doc/Magic-Values.md
  dims[u'dut_id'] = [unicode(common.id)]
  dims[u'dut_name'] = [unicode(common.hostname)]
  if common.HasField('hwid'):
    dims[u'hwid'] = [unicode(common.hwid)]
  if common.HasField('serial_number'):
    dims[u'serial_number'] = [unicode(common.serial_number)]
  if common.HasField('location'):
    dims[u'location'] = [u'%s-row%d-rack%d-host%d' % (
        common.location.lab.name,
        common.location.row,
        common.location.rack,
        common.location.host,
    )]
  if common.HasField('environment'):
    dims[u'environment'] = [unicode(
        common_pb2.Environment.Name(common.environment))]
  # TODO(Implement device locks by removing from inventory)

  dims.update(swarming.labels_to_dimensions(common.labels))
  return json.dumps(dims)


def main():
  """Prints swarming dimensions for the given dut."""
  parser = argparse.ArgumentParser(
      description='Print swarming dimensions for given DUT.')
  parser.add_argument(
      '--datadir',
      required=True,
      help='Path to the directory containing skylab inventory data.',
  )
  parser.add_argument(
      '--id',
      required=True,
      help='Inventory id of the DUT.',
  )
  parser.add_argument(
      '--environment',
      required=True,
      help='The inventory environment to query. One of prod, staging, skylab.',
  )
  opts = parser.parse_args()

  lab = text_manager.load_lab(opts.datadir, {opts.environment})
  for dut in lab.duts:
    if dut.common.id == opts.id:
      print(_dimensions_json(dut))
      exit(0)

  raise Exception(
      'Failed to find DUT with id %s in environment %s'
      % (opts.id, opts.environment))


if __name__ == '__main__':
  main()