venv/skylab_inventory/scripts/print_dut_state.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 state for a given DUT."""

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

from skylab_inventory import text_manager

import argparse
import collections
import json


def _state_json(dut):
  """Returns the state for swarming bot as json."""
  state = collections.defaultdict(list)
  for attribute in dut.common.attributes:
    state[attribute.key].append(attribute.value)
  return json.dumps(dict(state))


def main():
  """Prints swarming state for the given dut."""
  parser = argparse.ArgumentParser(
      description='Print swarming state 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.',
  )
  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(_state_json(dut))
      exit(0)

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


if __name__ == '__main__':
  main()