1#!/usr/bin/python 2# Copyright (c) 2012 The Chromium OS Authors. All rights reserved. 3# Use of this source code is governed by a BSD-style license that can be 4# found in the LICENSE file. 5 6import argparse 7import logging 8import sys 9import xmlrpclib 10 11import common 12 13from config import rpm_config 14from autotest_lib.client.common_lib import global_config 15from autotest_lib.client.common_lib.cros import retry 16 17RPM_FRONTEND_URI = global_config.global_config.get_config_value('CROS', 18 'rpm_frontend_uri', type=str, default='') 19RPM_CALL_TIMEOUT_MINS = rpm_config.getint('RPM_INFRASTRUCTURE', 20 'call_timeout_mins') 21 22POWERUNIT_HOSTNAME_KEY = 'powerunit_hostname' 23POWERUNIT_OUTLET_KEY = 'powerunit_outlet' 24HYDRA_HOSTNAME_KEY = 'hydra_hostname' 25 26 27class RemotePowerException(Exception): 28 """This is raised when we fail to set the state of the device's outlet.""" 29 pass 30 31 32def set_power(host, new_state, timeout_mins=RPM_CALL_TIMEOUT_MINS): 33 """Sends the power state change request to the RPM Infrastructure. 34 35 @param host: A CrosHost or ServoHost instance. 36 @param new_state: State we want to set the power outlet to. 37 """ 38 # servo V3 is handled differently from the rest. 39 # The best heuristic we have to determine servo V3 is the hostname. 40 if host.hostname.endswith('servo'): 41 args_tuple = (host.hostname, new_state) 42 else: 43 info = host.host_info_store.get() 44 try: 45 args_tuple = (host.hostname, 46 info.attributes[POWERUNIT_HOSTNAME_KEY], 47 info.attributes[POWERUNIT_OUTLET_KEY], 48 info.attributes.get(HYDRA_HOSTNAME_KEY), 49 new_state) 50 except KeyError as e: 51 raise RemotePowerException('Powerunit information not found. ' 52 'Missing: %s in data_info_store.' % e) 53 _set_power(args_tuple, timeout_mins) 54 55 56def _set_power(args_tuple, timeout_mins=RPM_CALL_TIMEOUT_MINS): 57 """Sends the power state change request to the RPM Infrastructure. 58 59 @param args_tuple: A args tuple for rpc call. See example below: 60 (hostname, powerunit_hostname, outlet, hydra_hostname, new_state) 61 """ 62 client = xmlrpclib.ServerProxy(RPM_FRONTEND_URI, 63 verbose=False, 64 allow_none=True) 65 timeout = None 66 result = None 67 endpoint = (client.set_power_via_poe if len(args_tuple) == 2 68 else client.set_power_via_rpm) 69 try: 70 timeout, result = retry.timeout(endpoint, 71 args=args_tuple, 72 timeout_sec=timeout_mins * 60, 73 default_result=False) 74 except Exception as e: 75 logging.exception(e) 76 raise RemotePowerException( 77 'Client call exception: ' + str(e)) 78 if timeout: 79 raise RemotePowerException( 80 'Call to RPM Infrastructure timed out.') 81 if not result: 82 error_msg = ('Failed to change outlet status for host: %s to ' 83 'state: %s.' % (args_tuple[0], args_tuple[-1])) 84 logging.error(error_msg) 85 raise RemotePowerException(error_msg) 86 87 88# This function will be removed once we try to move tests running in 89# the chaos lab to skylab. See crbug.com/863217. 90def set_power_afe(hostname, new_state, timeout_mins=RPM_CALL_TIMEOUT_MINS): 91 """Sends the power state change request to the RPM Infrastructure. 92 93 @param hostname: host who's power outlet we want to change. 94 @param new_state: State we want to set the power outlet to. 95 """ 96 client = xmlrpclib.ServerProxy(RPM_FRONTEND_URI, verbose=False) 97 timeout = None 98 result = None 99 try: 100 timeout, result = retry.timeout(client.queue_request, 101 args=(hostname, new_state), 102 timeout_sec=timeout_mins * 60, 103 default_result=False) 104 except Exception as e: 105 logging.exception(e) 106 raise RemotePowerException( 107 'Client call exception: ' + str(e)) 108 if timeout: 109 raise RemotePowerException( 110 'Call to RPM Infrastructure timed out.') 111 if not result: 112 error_msg = ('Failed to change outlet status for host: %s to ' 113 'state: %s.' % (hostname, new_state)) 114 logging.error(error_msg) 115 raise RemotePowerException(error_msg) 116 117 118def parse_options(): 119 """Parse the user supplied options.""" 120 parser = argparse.ArgumentParser() 121 parser.add_argument('-m', '--machine', dest='machine', required=True, 122 help='Machine hostname to change outlet state.') 123 parser.add_argument('-s', '--state', dest='state', required=True, 124 choices=['ON', 'OFF', 'CYCLE'], 125 help='Power state to set outlet: ON, OFF, CYCLE') 126 parser.add_argument('-p', '--powerunit_hostname', dest='p_hostname', 127 help='Powerunit hostname of the host.') 128 parser.add_argument('-o', '--outlet', dest='outlet', 129 help='Outlet of the host.') 130 parser.add_argument('-y', '--hydra_hostname', dest='hydra', default='', 131 help='Hydra hostname of the host.') 132 parser.add_argument('-d', '--disable_emails', dest='disable_emails', 133 help='Hours to suspend RPM email notifications.') 134 parser.add_argument('-e', '--enable_emails', dest='enable_emails', 135 action='store_true', 136 help='Resume RPM email notifications.') 137 return parser.parse_args() 138 139 140def main(): 141 """Entry point for rpm_client script.""" 142 options = parse_options() 143 if options.machine.endswith('servo'): 144 args_tuple = (options.machine, options.state) 145 elif not options.p_hostname or not options.outlet: 146 print "Powerunit hostname and outlet info are required for DUT." 147 return 148 else: 149 args_tuple = (options.machine, options.p_hostname, options.outlet, 150 options.hydra, options.state) 151 _set_power(args_tuple) 152 153 if options.disable_emails is not None: 154 client = xmlrpclib.ServerProxy(RPM_FRONTEND_URI, verbose=False) 155 client.suspend_emails(options.disable_emails) 156 if options.enable_emails: 157 client = xmlrpclib.ServerProxy(RPM_FRONTEND_URI, verbose=False) 158 client.resume_emails() 159 160 161if __name__ == "__main__": 162 sys.exit(main()) 163