Retry non-streaming gRPC calls

appengine/swarming/swarming_bot/bot_code/remote_client_grpc_test.py

Index: appengine/swarming/swarming_bot/bot_code/remote_client_grpc_test.py
diff --git a/appengine/swarming/swarming_bot/bot_code/remote_client_grpc_test.py b/appengine/swarming/swarming_bot/bot_code/remote_client_grpc_test.py
index 30085ee0c2b7c725594fd0142fdd1b65522e6eeb..7110c5bb6438a38a6ff6eb6733ecbb1b14213125 100755
--- a/appengine/swarming/swarming_bot/bot_code/remote_client_grpc_test.py
+++ b/appengine/swarming/swarming_bot/bot_code/remote_client_grpc_test.py
@@ -36,15 +36,44 @@ class FakeBotServiceStub(object):
     return self._testobj._handle_call('Poll', request)
 
 
+# If gRPC isn't successfully imported, this will be seen as a nonstandard
+# exception because it won't appear to be derived from Exception. This
+# only affects PyLint because the test will never be run if gRPC import
+# fails.
+# pylint: disable=W0710
+class FakeGrpcError(remote_client_grpc.grpc.RpcError):
+  """Duplicates a basic UNAVAILABLE error"""
+  def __init__(self, code):
+    self._code = code
+    super(FakeGrpcError, self).__init__('something terrible happened')
+
+  def code(self):
+    return self._code
+
+
 class TestRemoteClientGrpc(auto_stub.TestCase):
   def setUp(self):
     super(TestRemoteClientGrpc, self).setUp()
+    self._num_sleeps = 0
+    def fake_sleep(_time):
+      self._num_sleeps += 1
+    self.mock(time, 'sleep', fake_sleep)
     self._client = remote_client_grpc.RemoteClientGrpc('1.2.3.4:90')
     self._client._stub = FakeBotServiceStub(self)
     self._expected = []
+    self._error_codes = []
 
   def _handle_call(self, method, request):
     """This is called by FakeBotServiceStub to implement fake calls"""
+    if len(self._error_codes) > 0:
+      code, self._error_codes = self._error_codes[0], self._error_codes[1:]
+      raise FakeGrpcError(code)
+
+    if self._error_codes:
+      text = self._error_codes
+      self._error_codes = ''
+      raise FakeGrpcError(text)
+
     # Pop off the first item on the list
     self.assertTrue(len(self._expected) > 0)
     expected, self._expected = self._expected[0], self._expected[1:]
@@ -57,7 +86,7 @@ class TestRemoteClientGrpc(auto_stub.TestCase):
     return expected[2]
 
   def get_bot_attributes_dict(self):
-    """Gets the attributes of a basic bot"""
+    """Gets the attributes of a basic bot; must match next function"""
     return {
         'version': '123',
         'dimensions': {
@@ -125,6 +154,62 @@ class TestRemoteClientGrpc(auto_stub.TestCase):
         },
     })
 
+  def test_handshake_grpc_unavailable(self):
+    """Ensures that the handshake function sleeps after a gRPC error"""
+    msg_req = remote_client_grpc.swarming_bot_pb2.HandshakeRequest()
+    msg_req.attributes.CopyFrom(self.get_bot_attributes_proto())
+
+    # Create proto response
+    msg_rsp = remote_client_grpc.swarming_bot_pb2.HandshakeResponse()
+    msg_rsp.server_version = '101'
+    msg_rsp.bot_version = '102'
+    d1 = msg_rsp.bot_group_cfg.dimensions.add()
+    d1.name = 'mammal'
+    d1.values.extend(['kangaroo', 'emu'])
+
+    # Execute call and verify response
+    expected_call = ('Handshake', msg_req, msg_rsp)
+    self._expected.append(expected_call)
+    self._error_codes.append(remote_client_grpc.grpc.StatusCode.UNAVAILABLE)
+    self._error_codes.append(remote_client_grpc.grpc.StatusCode.UNAVAILABLE)
+    response = self._client.do_handshake(self.get_bot_attributes_dict())
+    self.assertEqual(self._num_sleeps, 2)
+    self.assertEqual(response, {
+        'server_version': u'101',
+        'bot_version': u'102',
+        'bot_group_cfg_version': u'',
+        'bot_group_cfg': {
+            'dimensions': {
+                u'mammal': [u'kangaroo', u'emu'],
+            },
+        },
+    })
+
+  def test_handshake_grpc_other_error(self):
+    """Ensures that the handshake function only catches UNAVAILABLE"""
+    self._error_codes.append(remote_client_grpc.grpc.StatusCode.UNAVAILABLE)
+    self._error_codes.append(remote_client_grpc.grpc.StatusCode.INTERNAL)
+    got_exception = None
+    try:
+      self._client.do_handshake(self.get_bot_attributes_dict())
+    except remote_client_grpc.grpc.RpcError as g:
+      got_exception = g
+    self.assertEqual(got_exception.code(),
+                     remote_client_grpc.grpc.StatusCode.INTERNAL)
+    self.assertEqual(self._num_sleeps, 1)
+
+  def test_handshake_grpc_too_many_errors(self):
+    """Ensures that the handshake function only catches UNAVAILABLE"""
+    self._error_codes.append(remote_client_grpc.grpc.StatusCode.UNAVAILABLE)
+    self._error_codes.append(remote_client_grpc.grpc.StatusCode.UNAVAILABLE)
+    self._error_codes.append(remote_client_grpc.grpc.StatusCode.UNAVAILABLE)
+    self.mock(remote_client_grpc, 'MAX_GRPC_ATTEMPTS', 2)
+    with self.assertRaises(remote_client_grpc.grpc.RpcError) as g:
+      self._client.do_handshake(self.get_bot_attributes_dict())
+    self.assertEqual(self._num_sleeps, 2)
+    self.assertEqual(g.exception.code(),
+                     remote_client_grpc.grpc.StatusCode.UNAVAILABLE)
+
   def test_poll_manifest(self):
     """Verifies that we can generate a reasonable manifest from a proto"""
     msg_req = remote_client_grpc.swarming_bot_pb2.PollRequest()