Fang Deng | 97951f0 | 2015-09-04 19:30:24 -0700 | [diff] [blame] | 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 | |
Fang Deng | 9f62788 | 2015-09-15 14:10:45 -0700 | [diff] [blame] | 42 | from autotest_lib.client.common_lib import global_config |
Fang Deng | 97951f0 | 2015-09-04 19:30:24 -0700 | [diff] [blame] | 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 | |
Fang Deng | 97951f0 | 2015-09-04 19:30:24 -0700 | [diff] [blame] | 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] |
Fang Deng | 9f62788 | 2015-09-15 14:10:45 -0700 | [diff] [blame] | 51 | DEFAULT_SERVER = global_config.global_config.get_config_value( |
| 52 | 'SERVER', 'hostname', default=None) |
Fang Deng | 97951f0 | 2015-09-04 19:30:24 -0700 | [diff] [blame] | 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') |
Fang Deng | 9f62788 | 2015-09-15 14:10:45 -0700 | [diff] [blame] | 147 | parser.add_argument('-s', '--server', type=str, dest='server', |
| 148 | default=DEFAULT_SERVER, |
Fang Deng | 97951f0 | 2015-09-04 19:30:24 -0700 | [diff] [blame] | 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(',')]) |