cli/rpc.py

   1 #
   2 # Copyright 2008 Google Inc. All Rights Reserved.
   3 #
   4
   5 import os
   6 from autotest_lib.frontend.afe import rpc_client_lib
   7 from autotest_lib.frontend.afe.json_rpc import proxy
   8 from autotest_lib.client.common_lib import global_config, utils
   9
  10 GLOBAL_CONFIG = global_config.global_config
  11 DEFAULT_SERVER = 'autotest'
  12 AFE_RPC_PATH = '/afe/server/rpc/'
  13 TKO_RPC_PATH = '/new_tko/server/rpc/'
  14
  15
  16 class AuthError(Exception):
  17     pass
  18
  19
  20 def get_autotest_server(web_server=None):
  21     if not web_server:
  22         if 'AUTOTEST_WEB' in os.environ:
  23             web_server = os.environ['AUTOTEST_WEB']
  24         else:
  25             web_server = 'http://' + GLOBAL_CONFIG.get_config_value(
  26                     'SERVER', 'hostname', default=DEFAULT_SERVER)
  27
  28     # if the name doesn't start with http://,
  29     # nonexistant hosts get an obscure error
  30     if not web_server.startswith('http://'):
  31         web_server = 'http://' + web_server
  32
  33     return web_server
  34
  35
  36 class rpc_comm(object):
  37     """Shared AFE/TKO RPC class stuff"""
  38     def __init__(self, web_server, rpc_path, username):
  39         self.username = username
  40         self.web_server = get_autotest_server(web_server)
  41         try:
  42             self.proxy = self._connect(rpc_path)
  43         except rpc_client_lib.AuthError, s:
  44             raise AuthError(s)
  45
  46
  47     def _connect(self, rpc_path):
  48         # This does not fail even if the address is wrong.
  49         # We need to wait for an actual RPC to fail
  50         headers = rpc_client_lib.authorization_headers(self.username,
  51                                                        self.web_server)
  52         rpc_server = self.web_server + rpc_path
  53         return rpc_client_lib.get_proxy(rpc_server, headers=headers)
  54
  55
  56     def run(self, op, *args, **data):
  57         if 'AUTOTEST_CLI_DEBUG' in os.environ:
  58             print self.web_server, op, args, data
  59         function = getattr(self.proxy, op)
  60         result = function(*args, **data)
  61         if 'AUTOTEST_CLI_DEBUG' in os.environ:
  62             print 'result:', result
  63         return result
  64
  65
  66 class afe_comm(rpc_comm):
  67     """Handles the AFE setup and communication through RPC"""
  68     def __init__(self, web_server=None, rpc_path=AFE_RPC_PATH, username=None):
  69         super(afe_comm, self).__init__(web_server, rpc_path, username)
  70
  71
  72 class tko_comm(rpc_comm):
  73     """Handles the TKO setup and communication through RPC"""
  74     def __init__(self, web_server=None, rpc_path=TKO_RPC_PATH, username=None):
  75         super(tko_comm, self).__init__(web_server, rpc_path, username)