| # Copyright 2014 The Chromium Authors. All rights reserved. | 
 | # Use of this source code is governed by a BSD-style license that can be | 
 | # found in the LICENSE file. | 
 |  | 
 | """Environment setup and teardown for remote devices.""" | 
 |  | 
 | import logging | 
 | import os | 
 | import random | 
 | import sys | 
 |  | 
 | from pylib import constants | 
 | from pylib.base import environment | 
 | from pylib.remote.device import appurify_sanitized | 
 | from pylib.remote.device import remote_device_helper | 
 |  | 
 | class RemoteDeviceEnvironment(environment.Environment): | 
 |   """An environment for running on remote devices.""" | 
 |  | 
 |   _ENV_KEY = 'env' | 
 |   _DEVICE_KEY = 'device' | 
 |  | 
 |   def __init__(self, args, error_func): | 
 |     """Constructor. | 
 |  | 
 |     Args: | 
 |       args: Command line arguments. | 
 |       error_func: error to show when using bad command line arguments. | 
 |     """ | 
 |     super(RemoteDeviceEnvironment, self).__init__() | 
 |  | 
 |     if args.api_key_file: | 
 |       with open(args.api_key_file) as api_key_file: | 
 |         self._api_key = api_key_file.read().strip() | 
 |     elif args.api_key: | 
 |       self._api_key = args.api_key | 
 |     else: | 
 |       error_func('Must set api key with --api-key or --api-key-file') | 
 |  | 
 |     if args.api_secret_file: | 
 |       with open(args.api_secret_file) as api_secret_file: | 
 |         self._api_secret = api_secret_file.read().strip() | 
 |     elif args.api_secret: | 
 |       self._api_secret = args.api_secret | 
 |     else: | 
 |       error_func('Must set api secret with --api-secret or --api-secret-file') | 
 |  | 
 |     if not args.api_protocol: | 
 |       error_func('Must set api protocol with --api-protocol. Example: http') | 
 |     self._api_protocol = args.api_protocol | 
 |  | 
 |     if not args.api_address: | 
 |       error_func('Must set api address with --api-address') | 
 |     self._api_address = args.api_address | 
 |  | 
 |     if not args.api_port: | 
 |       error_func('Must set api port with --api-port.') | 
 |     self._api_port = args.api_port | 
 |  | 
 |     self._access_token = '' | 
 |     self._results_path = args.results_path | 
 |     self._remote_device = args.remote_device | 
 |     self._remote_device_os = args.remote_device_os | 
 |     self._runner_package = args.runner_package | 
 |     self._runner_type = args.runner_type | 
 |     self._device = '' | 
 |     self._verbose_count = args.verbose_count | 
 |     self._device_type = args.device_type | 
 |     self._timeouts = { | 
 |         'queueing': 60 * 10, | 
 |         'installing': 60 * 10, | 
 |         'in-progress': 60 * 30, | 
 |         'unknown': 60 * 5 | 
 |     } | 
 |  | 
 |     if not args.trigger and not args.collect: | 
 |       self._trigger = True | 
 |       self._collect = True | 
 |     else: | 
 |       self._trigger = args.trigger | 
 |       self._collect = args.collect | 
 |  | 
 |   def SetUp(self): | 
 |     """Set up the test environment.""" | 
 |     os.environ['APPURIFY_API_PROTO'] = self._api_protocol | 
 |     os.environ['APPURIFY_API_HOST'] = self._api_address | 
 |     os.environ['APPURIFY_API_PORT'] = self._api_port | 
 |     self._GetAccessToken() | 
 |     if self._trigger: | 
 |       self._device = self._SelectDevice() | 
 |  | 
 |   def TearDown(self): | 
 |     """Teardown the test environment.""" | 
 |     self._RevokeAccessToken() | 
 |  | 
 |   def __enter__(self): | 
 |     """Set up the test run when used as a context manager.""" | 
 |     try: | 
 |       self.SetUp() | 
 |       return self | 
 |     except: | 
 |       self.__exit__(*sys.exc_info()) | 
 |       raise | 
 |  | 
 |   def __exit__(self, exc_type, exc_val, exc_tb): | 
 |     """Tears down the test run when used as a context manager.""" | 
 |     self.TearDown() | 
 |  | 
 |   def DumpTo(self, persisted_data): | 
 |     env_data = { | 
 |       self._DEVICE_KEY: self._device, | 
 |     } | 
 |     persisted_data[self._ENV_KEY] = env_data | 
 |  | 
 |   def LoadFrom(self, persisted_data): | 
 |     env_data = persisted_data[self._ENV_KEY] | 
 |     self._device = env_data[self._DEVICE_KEY] | 
 |  | 
 |   def _GetAccessToken(self): | 
 |     """Generates access token for remote device service.""" | 
 |     logging.info('Generating remote service access token') | 
 |     with appurify_sanitized.SanitizeLogging(self._verbose_count, | 
 |                                             logging.WARNING): | 
 |       access_token_results = appurify_sanitized.api.access_token_generate( | 
 |           self._api_key, self._api_secret) | 
 |     remote_device_helper.TestHttpResponse(access_token_results, | 
 |                                           'Unable to generate access token.') | 
 |     self._access_token = access_token_results.json()['response']['access_token'] | 
 |  | 
 |   def _RevokeAccessToken(self): | 
 |     """Destroys access token for remote device service.""" | 
 |     logging.info('Revoking remote service access token') | 
 |     with appurify_sanitized.SanitizeLogging(self._verbose_count, | 
 |                                             logging.WARNING): | 
 |       revoke_token_results = appurify_sanitized.api.access_token_revoke( | 
 |           self._access_token) | 
 |     remote_device_helper.TestHttpResponse(revoke_token_results, | 
 |                                           'Unable to revoke access token.') | 
 |  | 
 |   def _SelectDevice(self): | 
 |     """Select which device to use.""" | 
 |     logging.info('Finding device to run tests on.') | 
 |     with appurify_sanitized.SanitizeLogging(self._verbose_count, | 
 |                                             logging.WARNING): | 
 |       dev_list_res = appurify_sanitized.api.devices_list(self._access_token) | 
 |     remote_device_helper.TestHttpResponse(dev_list_res, | 
 |                                           'Unable to generate access token.') | 
 |     device_list = dev_list_res.json()['response'] | 
 |     random.shuffle(device_list) | 
 |     for device in device_list: | 
 |       if device['os_name'] != self._device_type: | 
 |         continue | 
 |       if self._remote_device and device['name'] != self._remote_device: | 
 |         continue | 
 |       if (self._remote_device_os | 
 |           and device['os_version'] != self._remote_device_os): | 
 |         continue | 
 |       if ((self._remote_device and self._remote_device_os) | 
 |           or device['available_devices_count']): | 
 |         logging.info('Found device: %s %s', | 
 |                      device['name'], device['os_version']) | 
 |         return device | 
 |     self._NoDeviceFound(device_list) | 
 |  | 
 |   def _PrintAvailableDevices(self, device_list): | 
 |     def compare_devices(a,b): | 
 |       for key in ('os_version', 'name'): | 
 |         c = cmp(a[key], b[key]) | 
 |         if c: | 
 |           return c | 
 |       return 0 | 
 |  | 
 |     logging.critical('Available %s Devices:', self._device_type) | 
 |     devices = (d for d in device_list if d['os_name'] == self._device_type) | 
 |     for d in sorted(devices, compare_devices): | 
 |       logging.critical('  %s %s', d['os_version'].ljust(7), d['name']) | 
 |  | 
 |   def _NoDeviceFound(self, device_list): | 
 |     self._PrintAvailableDevices(device_list) | 
 |     raise remote_device_helper.RemoteDeviceError('No device found.') | 
 |  | 
 |   @property | 
 |   def collect(self): | 
 |     return self._collect | 
 |  | 
 |   @property | 
 |   def device_type_id(self): | 
 |     return self._device['device_type_id'] | 
 |  | 
 |   @property | 
 |   def only_output_failures(self): | 
 |     # TODO(jbudorick): Remove this once b/18981674 is fixed. | 
 |     return True | 
 |  | 
 |   @property | 
 |   def results_path(self): | 
 |     return self._results_path | 
 |  | 
 |   @property | 
 |   def runner_package(self): | 
 |     return self._runner_package | 
 |  | 
 |   @property | 
 |   def runner_type(self): | 
 |     return self._runner_type | 
 |  | 
 |   @property | 
 |   def timeouts(self): | 
 |     return self._timeouts | 
 |  | 
 |   @property | 
 |   def token(self): | 
 |     return self._access_token | 
 |  | 
 |   @property | 
 |   def trigger(self): | 
 |     return self._trigger | 
 |  | 
 |   @property | 
 |   def verbose_count(self): | 
 |     return self._verbose_count | 
 |  | 
 |   @property | 
 |   def device_type(self): | 
 |     return self._device_type |