Add timeouts and unit tests to remote_client_grpc.py

appengine/swarming/swarming_bot/bot_code/remote_client_grpc.py

 # Copyright 2016 The LUCI Authors. All rights reserved.
 # Use of this source code is governed under the Apache License, Version 2.0
 # that can be found in the LICENSE file.
 
 # This is a reimplementation of RemoteClientNative but it uses (will use)
 # a gRPC method to communicate with a server instead of REST.
 
 import json
 import logging
 
 import grpc
 import google.protobuf.json_format
 from proto_bot import swarming_bot_pb2
 from remote_client_errors import InternalError
 
 
+# How long to wait for a response from the server. Keeping the same as
+# the equivalent in remote_client.py for now.
+NET_CONNECTION_TIMEOUT_SEC = 5*60
+
+
 class RemoteClientGrpc(object):
   """RemoteClientGrpc knows how to make calls via gRPC.
   """
 
   def __init__(self, server):
     logging.info('Communicating with host %s via gRPC', server)
     self._server = server
     self._channel = grpc.insecure_channel(server)
     self._stub = swarming_bot_pb2.BotServiceStub(self._channel)
     self._log_is_asleep = False
@@ skipping 28 matching lines @@
       request.output_chunk.data = stdout_and_chunk[0]
       request.output_chunk.offset = stdout_and_chunk[1]
     if exit_code != None:
       request.exit_status.code = exit_code
 
     # Insert everything else. Note that the b64-encoded strings in the dict
     # are automatically decoded by ParseDict.
     google.protobuf.json_format.ParseDict(params, request)
 
     # Perform update
-    response = self._stub.TaskUpdate(request)
+    response = self._stub.TaskUpdate(request,
+                                     timeout=NET_CONNECTION_TIMEOUT_SEC)
     logging.debug('post_task_update() = %s', request)
     if response.error:
       raise InternalError(response.error)
     return not response.must_stop
 
   def post_task_error(self, task_id, bot_id, message):
     request = swarming_bot_pb2.TaskErrorRequest()
     request.bot_id = bot_id
     request.task_id = task_id
     request.msg = message
     logging.error('post_task_error() = %s', request)
 
-    response = self._stub.TaskError(request)
+    response = self._stub.TaskError(request, timeout=NET_CONNECTION_TIMEOUT_SEC)
     return response.ok
 
   def _attributes_json_to_proto(self, json_attr, msg):
     msg.version = json_attr['version']
     for k, values in sorted(json_attr['dimensions'].iteritems()):
       pair = msg.dimensions.add()
       pair.name = k
       pair.values.extend(values)
     create_state_proto(json_attr['state'], msg.state)
 
   def do_handshake(self, attributes):
     request = swarming_bot_pb2.HandshakeRequest()
     self._attributes_json_to_proto(attributes, request.attributes)
-    response = self._stub.Handshake(request)
+    response = self._stub.Handshake(request, timeout=NET_CONNECTION_TIMEOUT_SEC)
     resp = {
         'server_version': response.server_version,
         'bot_version': response.bot_version,
         'bot_group_cfg_version': response.bot_group_cfg_version,
         'bot_group_cfg': {
             'dimensions': {
                 d.name: d.values for d in response.bot_group_cfg.dimensions
             },
         },
     }
     logging.info('Completed handshake: %s', resp)
     return resp
 
   def poll(self, attributes):
     request = swarming_bot_pb2.PollRequest()
     self._attributes_json_to_proto(attributes, request.attributes)
     # TODO(aludwin): gRPC-specific exception handling
-    response = self._stub.Poll(request)
+    response = self._stub.Poll(request, timeout=NET_CONNECTION_TIMEOUT_SEC)
 
     if response.cmd == swarming_bot_pb2.PollResponse.UPDATE:
       return 'update', response.version
 
     if response.cmd == swarming_bot_pb2.PollResponse.SLEEP:
       if not self._log_is_asleep:
         logging.info('Going to sleep')
         self._log_is_asleep = True
       return 'sleep', response.sleep_time
 
@@ skipping 39 matching lines @@
       self._log_is_asleep = False
       return 'run', manifest
 
     raise ValueError('Unknown command in response: %s' % response)
 
   def get_bot_code(self, new_zip_fn, bot_version, _bot_id):
     # TODO(aludwin): exception handling, pass bot_id
     logging.info('Updating to version: %s', bot_version)
     request = swarming_bot_pb2.BotUpdateRequest()
     request.bot_version = bot_version
-    response = self._stub.BotUpdate(request)
+    response = self._stub.BotUpdate(request, timeout=NET_CONNECTION_TIMEOUT_SEC)
     with open(new_zip_fn, 'wb') as f:
       f.write(response.bot_code)
 
   def ping(self):
     pass
 
 
 def create_state_proto(state_dict, message):
   """ Constructs a State message out of a state dict.
 
@@ skipping 27 matching lines @@
 def insert_dict_as_submessage(message, keyname, value):
   """Encodes a dict as a Protobuf message.
 
   The keyname for the message field is passed in to simplify the creation
   of the submessage in the first place - you need to say getattr and not
   simply refer to message.keyname since the former actually creates the
   submessage while the latter does not.
   """
   sub_msg = getattr(message, keyname)
   google.protobuf.json_format.Parse(json.dumps(value), sub_msg)