Refactor all gRPC proxy code into a single class.

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 7110c5bb6438a38a6ff6eb6733ecbb1b14213125..76f4aa942c483d52ee6ce219c9b02a30466dfcc1 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
@@ -13,42 +13,15 @@ import test_env_bot_code
 test_env_bot_code.setup_test_env()
 
 from depot_tools import auto_stub
+import remote_client_grpc
 
-try:
-  import remote_client_grpc
-except ImportError as e:
-  print('Could not import gRPC remote client, likely due to missing grpc '
-        'library. Skipping tests.')
-  sys.exit(0)
 
-
-class FakeBotServiceStub(object):
+class FakeGrpcProxy(object):
   def __init__(self, testobj):
     self._testobj = testobj
 
-  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)
-
-
-# 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
+  def call_unary(self, name, request):
+    return self._testobj._handle_call(name, request)
 
 
 class TestRemoteClientGrpc(auto_stub.TestCase):
@@ -58,22 +31,13 @@ class TestRemoteClientGrpc(auto_stub.TestCase):
     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._client = remote_client_grpc.RemoteClientGrpc('1.2.3.4:90',
+                                                       FakeGrpcProxy(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)
-
+    """This is called by FakeGrpcProxy to implement fake calls"""
     # Pop off the first item on the list
     self.assertTrue(len(self._expected) > 0)
     expected, self._expected = self._expected[0], self._expected[1:]
@@ -154,62 +118,6 @@ 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()