1 #! /usr/bin/python 2 3 # Copyright (c) 2014 The Chromium OS Authors. All rights reserved. 4 # Use of this source code is governed by a BSD-style license that can be 5 # found in the LICENSE file. 6 7 """ 8 Manage power unit information for autotest hosts. 9 10 We store rpm hostname, outlet, hydra information for a host in cautotest 11 as host attributes. This tool allows you to add/modify/view/backup 12 rpm attributes for hosts. 13 14 * Add/Modify power unit attributes: 15 Step 1: create csv: 16 Put attributes in a csv file, e.g. mapping.csv. 17 Each line in mapping.csv consists of 18 device_hostname, powerunit_hostname, powerunit_outlet, hydra_hostname, 19 seperated by comma. For example 20 21 chromeos-rack2-host1,chromeos-rack2-rpm1,.A1,chromeos-197-hydra1.mtv, 22 chromeos-rack2-host2,chromeos-rack2-rpm1,.A2,chromeos-197-hydra1.mtv, 23 24 Step 2: run 25 ./manage_powerunit_info.py upload --csv mapping_file.csv 26 27 * View power unit attributes: 28 ./manage_powerunit_info.py list 29 -m "chromeos-rack2-host1,chromeos-rack2-host2" 30 31 * Backup existing attributes for all hosts to a csv file: 32 ./manage_powerunit_info.py backup --csv backup.csv 33 """ 34 import argparse 35 import csv 36 import logging 37 import os 38 import sys 39 40 import common 41 42 from autotest_lib.client.common_lib import global_config 43 from autotest_lib.server.cros.dynamic_suite import frontend_wrappers 44 from autotest_lib.site_utils.rpm_control_system import utils as rpm_utils 45 46 47 # The host attribute key name for get rpm hostname. 48 POWERUNIT_KEYS = [rpm_utils.POWERUNIT_HOSTNAME_KEY, 49 rpm_utils.POWERUNIT_OUTLET_KEY, 50 rpm_utils.HYDRA_HOSTNAME_KEY] 51 DEFAULT_SERVER = global_config.global_config.get_config_value( 52 'SERVER', 'hostname', default=None) 53 54 55 def add_powerunit_info_to_host(afe, device, keyvals): 56 """Add keyvals to the host's attributes in AFE. 57 58 @param afe: AFE server to talk to. 59 @param device: the device hostname, e.g. 'chromeos1-rack1-host1' 60 @param keyvals: A dictionary where keys are the values in POWERUNIT_KEYS. 61 These are the power unit info about the devcie that we 62 are going to insert to AFE as host attributes. 63 """ 64 if not afe.get_hosts(hostname=device): 65 logging.debug('No host named %s', device) 66 return 67 68 logging.info('Adding host attribues to %s: %s', device, keyvals) 69 for key, val in keyvals.iteritems(): 70 afe.set_host_attribute(key, val, hostname=device) 71 72 73 def add_from_csv(afe, csv_file): 74 """Read power unit information from csv and add to host attributes. 75 76 @param afe: AFE server to talk to. 77 @param csv_file: A csv file, each line consists of device_hostname, 78 powerunit_hostname powerunit_outlet, hydra_hostname 79 separated by comma. 80 """ 81 with open(csv_file) as f: 82 reader = csv.reader(f, delimiter=',') 83 for row in reader: 84 device = row[0].strip() 85 hydra = row[3].strip() 86 if not hydra: 87 hydra = None 88 keyvals = dict(zip( 89 POWERUNIT_KEYS, 90 [row[1].strip(), row[2].strip(), hydra])) 91 add_powerunit_info_to_host(afe, device, keyvals) 92 93 94 def dump_to_csv(afe, csv_file): 95 """Dump power unit info of all hosts to a csv file. 96 97 @param afe: AFE server to talk to. 98 @param csv_file: A file to store the power unit information. 99 100 """ 101 logging.info('Back up host attribues to %s', csv_file) 102 with open(csv_file, 'w') as f: 103 hosts = afe.get_hosts() 104 for h in hosts: 105 logging.info('Proccessing %s', h.hostname) 106 f.write(h.hostname + ',') 107 for key in POWERUNIT_KEYS: 108 f.write(h.attributes.get(key, '') + ',') 109 f.write('\n') 110 111 112 def list_powerunit_info(afe, devices): 113 """List power unit info for a list of hosts. 114 115 @param afe: AFE server to talk to. 116 @param devices: a list of device hostnames. 117 """ 118 hosts = afe.get_hosts(hostname__in = devices) 119 if not hosts: 120 logging.error('No host found.') 121 for h in hosts: 122 info = h.hostname + ',' 123 for key in POWERUNIT_KEYS: 124 info += h.attributes.get(key, '') + ',' 125 print info 126 127 128 def parse_options(): 129 """Parse options""" 130 parser = argparse.ArgumentParser( 131 description=__doc__, 132 formatter_class=argparse.RawDescriptionHelpFormatter) 133 action_help = ( 134 'upload: read rpm attributes from csv file and set the attributes. ' 135 'list: list current attributes for a list of hosts. ' 136 'backup: dump existing rpm attributes to a csv file (for backup).') 137 parser.add_argument( 138 'action', choices=('upload', 'list', 'backup'), help=action_help) 139 parser.add_argument('-f', '--csv_file', type=str, dest='csv_file', 140 help='A path to a csv file. When upload, each line ' 141 'should consist of device_name, powerunit_hostname, ' 142 'powerunit_outlet, hydra_hostname, separated ' 143 'by comma. When dump, the file will be generated.') 144 parser.add_argument('-m', type=str, dest='hostnames', default='', 145 help='A list of machine hostnames seperated by comma, ' 146 'applicable to "list" command') 147 parser.add_argument('-s', '--server', type=str, dest='server', 148 default=DEFAULT_SERVER, 149 help='AFE server that the script will be talking to. ' 150 'If not speicified, will default to using the ' 151 'server in global_config.ini') 152 options = parser.parse_args() 153 if options.action == 'upload' or options.action =='backup': 154 if not options.csv_file: 155 logging.error('Please specifiy a file with -f/--csv') 156 sys.exit(1) 157 file_exists = os.path.exists(options.csv_file) 158 if options.action == 'upload' and not file_exists: 159 logging.error('%s is not a valid file.', options.csv_file) 160 sys.exit(1) 161 if options.action == 'backup' and file_exists: 162 logging.error('%s already exists.', options.csv_file) 163 sys.exit(1) 164 if options.action == 'list' and not options.hostnames: 165 logging.error('Please specify hostnames with -m') 166 sys.exit(1) 167 return options 168 169 170 if __name__ == '__main__': 171 logging.basicConfig(level=logging.DEBUG) 172 options = parse_options() 173 afe = frontend_wrappers.RetryingAFE(timeout_min=5, delay_sec=10, 174 server=options.server) 175 logging.info('Connected to %s', afe.server) 176 if options.action =='backup': 177 dump_to_csv(afe, options.csv_file) 178 elif options.action == 'upload': 179 confirm_msg = ('Upload rpm mapping from %s, are you sure?' 180 % options.csv_file) 181 confirm = raw_input("%s (y/N) " % confirm_msg).lower() == 'y' 182 if confirm: 183 add_from_csv(afe, options.csv_file) 184 elif options.action == 'list': 185 list_powerunit_info(afe, [h.strip() for h in options.hostnames.split(',')]) 186
