Retry non-streaming gRPC calls

appengine/swarming/swarming_bot/bot_code/remote_client_grpc_test.py

 #!/usr/bin/env python
 # 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.
 
 import logging
 import sys
 import threading
 import time
 import unittest
@@ skipping 18 matching lines @@
   def TaskUpdate(self, request, **_kwargs):
     return self._testobj._handle_call('TaskUpdate', request)
 
   def Handshake(self, request, **_kwargs):
     return self._testobj._handle_call('Handshake', request)
 
   def Poll(self, request, **_kwargs):
     return self._testobj._handle_call('Poll', request)
 
 
+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
+

[M-A Ruel, 2016/12/21 16:31:44]

add empty line here

[aludwin, 2016/12/21 16:50:35]

Done.

 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:]
     # Each element of the "expected" array should be a 3-tuple:
     #    * The name of the method (eg 'TaskUpdate')
     #    * The proto request
     #    * The proto response
     self.assertEqual(method, expected[0])
     self.assertEqual(request, expected[1])
     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': {
             'mammal': ['ferrett', 'wombat'],
             'pool': ['dead'],
         },
         'state': {
             'audio': ['Rachmaninov', 'Stravinsky'],
             'cost_usd_hour': 3.141,
             'cpu': 'Intel 8080',
@@ skipping 47 matching lines @@
         '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_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)
+    got_exception = None
+    try:

[M-A Ruel, 2016/12/21 16:31:44]

FYI, you can do instead:

with self.asserRaises(remote_client_grpc.grpc.RpcError) as e:
  self._client.do_handshake(self.get_bot_attributes_dict())
self.assertEqual(self._num_sleeps, 2)
self.assertEqual(e.code(),
                 remote_client_grpc.grpc.StatusCode.UNAVAILABLE)

[aludwin, 2016/12/21 16:50:35]

Done.

+      self._client.do_handshake(self.get_bot_attributes_dict())
+    except remote_client_grpc.grpc.RpcError as g:
+      got_exception = g
+    self.assertEqual(self._num_sleeps, 2)
+    self.assertEqual(got_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()
     msg_req.attributes.CopyFrom(self.get_bot_attributes_proto())
 
     # Create dict response
     dict_rsp = {
         'bot_id': u'baby_bot',
         'command': None,
         'dimensions': {
@@ skipping 107 matching lines @@
     self._expected.append(expected_call)
     response = self._client.post_task_update('abc123', 'baby_bot', params,
                                              ['abc', 7], -5)
     self.assertTrue(response)
 
 if __name__ == '__main__':
   logging.basicConfig(
       level=logging.DEBUG if '-v' in sys.argv else logging.CRITICAL)
   unittest.TestCase.maxDiff = None
   unittest.main()