blob: f8f27d749fbe9331f3067c50bcbfbbbadcf25a0c [file] [log] [blame]
Fang Deng97951f02015-09-04 19:30:24 -07001#! /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"""
8Manage 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"""
34import argparse
35import csv
36import logging
37import os
38import sys
39
40import common
41
Fang Deng9f627882015-09-15 14:10:45 -070042from autotest_lib.client.common_lib import global_config
Fang Deng97951f02015-09-04 19:30:24 -070043from autotest_lib.server.cros.dynamic_suite import frontend_wrappers
44from autotest_lib.site_utils.rpm_control_system import utils as rpm_utils
45
46
Fang Deng97951f02015-09-04 19:30:24 -070047# The host attribute key name for get rpm hostname.
48POWERUNIT_KEYS = [rpm_utils.POWERUNIT_HOSTNAME_KEY,
49 rpm_utils.POWERUNIT_OUTLET_KEY,
50 rpm_utils.HYDRA_HOSTNAME_KEY]
Fang Deng9f627882015-09-15 14:10:45 -070051DEFAULT_SERVER = global_config.global_config.get_config_value(
52 'SERVER', 'hostname', default=None)
Fang Deng97951f02015-09-04 19:30:24 -070053
54
55def 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
73def 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
94def 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
112def 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
128def 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 Deng9f627882015-09-15 14:10:45 -0700147 parser.add_argument('-s', '--server', type=str, dest='server',
148 default=DEFAULT_SERVER,
Fang Deng97951f02015-09-04 19:30:24 -0700149 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
170if __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(',')])