Refactor all gRPC proxy code into a single class.

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
 
 import test_env_bot_code
 test_env_bot_code.setup_test_env()
 
 from depot_tools import auto_stub
-
-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)
+import remote_client_grpc
 
 
-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):
   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._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:]
     # 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]
@@ skipping 60 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)
-    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()
     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()