1#!/usr/bin/python2
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
17try:
18    from chromite.lib import metrics
19except ImportError:
20    from autotest_lib.client.bin.utils import metrics_mock as metrics
21
22RPM_FRONTEND_URI = global_config.global_config.get_config_value('CROS',
23        'rpm_frontend_uri', type=str, default='')
24RPM_CALL_TIMEOUT_MINS = rpm_config.getint('RPM_INFRASTRUCTURE',
25                                          'call_timeout_mins')
26
27POWERUNIT_HOSTNAME_KEY = 'powerunit_hostname'
28POWERUNIT_OUTLET_KEY = 'powerunit_outlet'
29HYDRA_HOSTNAME_KEY = 'hydra_hostname'
30
31
32class RemotePowerException(Exception):
33    """This is raised when we fail to set the state of the device's outlet."""
34    pass
35
36
37def set_power(host, new_state, timeout_mins=RPM_CALL_TIMEOUT_MINS):
38    """Sends the power state change request to the RPM Infrastructure.
39
40    @param host: A CrosHost or ServoHost instance.
41    @param new_state: State we want to set the power outlet to.
42    """
43    # servo V3 is handled differently from the rest.
44    # The best heuristic we have to determine servo V3 is the hostname.
45    if host.hostname.endswith('servo'):
46        args_tuple = (host.hostname, new_state)
47    else:
48        info = host.host_info_store.get()
49        try:
50            args_tuple = (host.hostname,
51                          info.attributes[POWERUNIT_HOSTNAME_KEY],
52                          info.attributes[POWERUNIT_OUTLET_KEY],
53                          info.attributes.get(HYDRA_HOSTNAME_KEY),
54                          new_state)
55        except KeyError as e:
56            logging.warning('Powerunit information not found. Missing:'
57                            ' %s in host_info_store.', e)
58            raise RemotePowerException('Remote power control is not applicable'
59                                       ' for %s, it could be either RPM is not'
60                                       ' supported on the rack or powerunit'
61                                       ' attributes is not configured in'
62                                       ' inventory.' % host.hostname)
63    _set_power(args_tuple, timeout_mins)
64
65
66def _set_power(args_tuple, timeout_mins=RPM_CALL_TIMEOUT_MINS):
67    """Sends the power state change request to the RPM Infrastructure.
68
69    @param args_tuple: A args tuple for rpc call. See example below:
70        (hostname, powerunit_hostname, outlet, hydra_hostname, new_state)
71    """
72    client = xmlrpclib.ServerProxy(RPM_FRONTEND_URI,
73                                   verbose=False,
74                                   allow_none=True)
75    timeout = None
76    result = None
77    endpoint = (client.set_power_via_poe if len(args_tuple) == 2
78                else client.set_power_via_rpm)
79    try:
80        timeout, result = retry.timeout(endpoint,
81                                        args=args_tuple,
82                                        timeout_sec=timeout_mins * 60,
83                                        default_result=False)
84    except Exception as e:
85        logging.exception(e)
86        raise RemotePowerException(
87                'Client call exception (%s): %s' % (RPM_FRONTEND_URI, e))
88    if timeout:
89        raise RemotePowerException(
90                'Call to RPM Infrastructure timed out (%s).' % RPM_FRONTEND_URI)
91    if not result:
92        error_msg = ('Failed to change outlet status for host: %s to '
93                     'state: %s.' % (args_tuple[0], args_tuple[-1]))
94        logging.error(error_msg)
95        if len(args_tuple) > 2:
96            # Collect failure metrics if we set power via rpm.
97            _send_rpm_failure_metrics(args_tuple[0], args_tuple[1],
98                                      args_tuple[2])
99        raise RemotePowerException(error_msg)
100
101
102def _send_rpm_failure_metrics(hostname, rpm_host, outlet):
103    metrics_fields = {
104            'hostname': hostname,
105            'rpm_host': rpm_host,
106            'outlet': outlet
107    }
108    metrics.Counter('chromeos/autotest/rpm/rpm_failure2').increment(
109            fields=metrics_fields)
110
111
112def parse_options():
113    """Parse the user supplied options."""
114    parser = argparse.ArgumentParser()
115    parser.add_argument('-m', '--machine', dest='machine', required=True,
116                        help='Machine hostname to change outlet state.')
117    parser.add_argument('-s', '--state', dest='state', required=True,
118                        choices=['ON', 'OFF', 'CYCLE'],
119                        help='Power state to set outlet: ON, OFF, CYCLE')
120    parser.add_argument('-p', '--powerunit_hostname', dest='p_hostname',
121                        help='Powerunit hostname of the host.')
122    parser.add_argument('-o', '--outlet', dest='outlet',
123                        help='Outlet of the host.')
124    parser.add_argument('-y', '--hydra_hostname', dest='hydra', default='',
125                        help='Hydra hostname of the host.')
126    parser.add_argument('-d', '--disable_emails', dest='disable_emails',
127                        help='Hours to suspend RPM email notifications.')
128    parser.add_argument('-e', '--enable_emails', dest='enable_emails',
129                        action='store_true',
130                        help='Resume RPM email notifications.')
131    return parser.parse_args()
132
133
134def main():
135    """Entry point for rpm_client script."""
136    options = parse_options()
137    if options.machine.endswith('servo'):
138        args_tuple = (options.machine, options.state)
139    elif not options.p_hostname or not options.outlet:
140        print "Powerunit hostname and outlet info are required for DUT."
141        return
142    else:
143        args_tuple = (options.machine, options.p_hostname, options.outlet,
144                      options.hydra, options.state)
145    _set_power(args_tuple)
146
147    if options.disable_emails is not None:
148        client = xmlrpclib.ServerProxy(RPM_FRONTEND_URI, verbose=False)
149        client.suspend_emails(options.disable_emails)
150    if options.enable_emails:
151        client = xmlrpclib.ServerProxy(RPM_FRONTEND_URI, verbose=False)
152        client.resume_emails()
153
154
155if __name__ == "__main__":
156    sys.exit(main())
157