Add GatewayTimeout exception to pubsub subscription pull.

daspecster · daspecster · commit 2a078e01a2bc · 2016-10-30T16:11:30.000-04:00
diff --git a/core/google/cloud/_testing.py b/core/google/cloud/_testing.py
@@ -77,6 +77,10 @@ def _make_grpc_failed_precondition(self):
         from grpc import StatusCode
         return self._make_grpc_error(StatusCode.FAILED_PRECONDITION)
 
+    def _make_grpc_deadline_exceeded(self):
+        from grpc import StatusCode
+        return self._make_grpc_error(StatusCode.DEADLINE_EXCEEDED)
+
 
 class _GAXPageIterator(object):
 
diff --git a/pubsub/google/cloud/pubsub/_gax.py b/pubsub/google/cloud/pubsub/_gax.py
@@ -31,6 +31,7 @@
 from google.cloud._helpers import _to_bytes
 from google.cloud._helpers import _pb_timestamp_to_rfc3339
 from google.cloud.exceptions import Conflict
+from google.cloud.exceptions import GatewayTimeout
 from google.cloud.exceptions import NotFound
 from google.cloud.iterator import GAXIterator
 from google.cloud.pubsub._helpers import subscription_name_from_path
@@ -413,6 +414,8 @@ def subscription_pull(self, subscription_path, return_immediately=False,
         except GaxError as exc:
             if exc_to_code(exc.cause) == StatusCode.NOT_FOUND:
                 raise NotFound(subscription_path)
+            if exc_to_code(exc.cause) == StatusCode.DEADLINE_EXCEEDED:
+                raise GatewayTimeout(subscription_path)
             raise
         return [_received_message_pb_to_mapping(rmpb)
                 for rmpb in response_pb.received_messages]
diff --git a/pubsub/unit_tests/test__gax.py b/pubsub/unit_tests/test__gax.py
@@ -762,6 +762,14 @@ def test_subscription_pull_defaults_error(self):
         self.assertFalse(return_immediately)
         self.assertIsNone(options)
 
+    def test_subscription_pull_deadline_exceeded(self):
+        from google.cloud.exceptions import GatewayTimeout
+        gax_api = _GAXSubscriberAPI(_deadline_exceeded_gax_error=True)
+        api = self._makeOne(gax_api)
+
+        with self.assertRaises(GatewayTimeout):
+            api.subscription_pull(self.SUB_PATH)
+
     def test_subscription_acknowledge_hit(self):
         ACK_ID1 = 'DEADBEEF'
         ACK_ID2 = 'BEADCAFE'
@@ -1035,6 +1043,7 @@ class _GAXSubscriberAPI(_GAXBaseAPI):
     _modify_push_config_ok = False
     _acknowledge_ok = False
     _modify_ack_deadline_ok = False
+    _deadline_exceeded_gax_error = False
 
     def list_subscriptions(self, project, page_size, options=None):
         self._list_subscriptions_called_with = (project, page_size, options)
@@ -1084,6 +1093,9 @@ def pull(self, name, max_messages, return_immediately, options=None):
             name, max_messages, return_immediately, options)
         if self._random_gax_error:
             raise GaxError('error')
+        if self._deadline_exceeded_gax_error:
+            raise GaxError('deadline exceeded',
+                           self._make_grpc_deadline_exceeded())
         try:
             return self._pull_response
         except AttributeError:
