catatester/devicemap/devicemap.py

"""CLI to query bot dimensions for the device(s) of a given bot.

Example usage:
  python devicemap.py home-swarming-bot-host-3--7
  (writes bot dimensions & state as python dict to stdout)
"""

import argparse
import json
import logging
import os
import re
import yaml

# This code runs on catatester host and doesn't have whole code tree, hence
# do relative import.
try:
  from catatester.devicemap import devicemap_lib
except ImportError:
  import devicemap_lib # pylint: disable=relative-import


_BOTNAME_PATTERN = r'(.*)--(\d*)'


def _ParseArgs():
  """Parses command-line input."""
  parser = argparse.ArgumentParser()

  parser.add_argument(
    'botname',
    help='Name of the bot, in the format <hostname>--<bot #>',
  )

  parser.add_argument(
    '--output_file',
    help='Absolute path to the output file to write dimensions (as json).'
  )

  # Argument Validations
  args = parser.parse_args()
  match = re.match(_BOTNAME_PATTERN, args.botname)
  if not match:
    parser.error('invalid botname')
  if not int(match.group(2)) > 0:
    parser.error('bot # must be > 0')
  logging.info('Running with: %s', args)
  return args


def main():
  flags = _ParseArgs()
  match = re.match(_BOTNAME_PATTERN, flags.botname)
  hostname = match.group(1)
  bot_num = int(match.group(2))

  configs = devicemap_lib.Configs()
  bot_config = configs.GetBotConfig(hostname, bot_num)
  print(bot_config)
  if flags.output_file:
    with open(flags.output_file, 'w') as f:
      f.write(json.dumps(bot_config.as_dict))


if __name__ == '__main__':
  main()