feat: add timeout parameter to Batch interface to match google-cloud-core

storage/google/cloud/storage/batch.py

@@ -150,7 +150,7 @@ def __init__(self, client):
         self._requests = []
         self._target_objects = []
 
-    def _do_request(self, method, url, headers, data, target_object):
+    def _do_request(self, method, url, headers, data, target_object, timeout=None):
         """Override Connection:  defer actual HTTP request.
 
         Only allow up to ``_MAX_BATCH_SIZE`` requests to be deferred.
@@ -173,6 +173,12 @@ def _do_request(self, method, url, headers, data, target_object):
             connection. Here we defer an HTTP request and complete
             initialization of the object at a later time.
 
+        :type timeout: float or tuple
+        :param timeout: (optional) The amount of time, in seconds, to wait
+            for the server response. By default, the method waits indefinitely.
+            Can also be passed as a tuple (connect_timeout, read_timeout).
+            See :meth:`requests.Session.request` documentation for details.
+
         :rtype: tuple of ``response`` (a dictionary of sorts)
                 and ``content`` (a string).
         :returns: The HTTP response object and the content of the response.
@@ -181,7 +187,7 @@ def _do_request(self, method, url, headers, data, target_object):
             raise ValueError(
                 "Too many deferred requests (max %d)" % self._MAX_BATCH_SIZE
             )
-        self._requests.append((method, url, headers, data))
+        self._requests.append((method, url, headers, data, timeout))
         result = _FutureDict()
         self._target_objects.append(target_object)
         if target_object is not None:
@@ -200,9 +206,12 @@ def _prepare_batch_request(self):
 
         multi = MIMEMultipart()
 
-        for method, uri, headers, body in self._requests:
+        # Use timeout of last request, default to None (indefinite)
+        timeout = None
+        for method, uri, headers, body, _timeout in self._requests:
             subrequest = MIMEApplicationHTTP(method, uri, headers, body)
             multi.attach(subrequest)
+            timeout = _timeout
 
         # The `email` package expects to deal with "native" strings
         if six.PY3:  # pragma: NO COVER  Python3
@@ -215,7 +224,7 @@ def _prepare_batch_request(self):
 
         # Strip off redundant header text
         _, body = payload.split("\n\n", 1)
-        return dict(multi._headers), body
+        return dict(multi._headers), body, timeout
 
     def _finish_futures(self, responses):
         """Apply all the batch responses to the futures created.
@@ -230,7 +239,7 @@ def _finish_futures(self, responses):
         # until all futures have been populated.
         exception_args = None
 
-        if len(self._target_objects) != len(responses):
+        if len(self._target_objects) != len(responses):  # pragma: NO COVER
             raise ValueError("Expected a response for every request.")
 
         for target_object, subresponse in zip(self._target_objects, responses):
@@ -251,15 +260,15 @@ def finish(self):
         :rtype: list of tuples
         :returns: one ``(headers, payload)`` tuple per deferred request.
         """
-        headers, body = self._prepare_batch_request()
+        headers, body, timeout = self._prepare_batch_request()
 
         url = "%s/batch/storage/v1" % self.API_BASE_URL
 
         # Use the private ``_base_connection`` rather than the property
         # ``_connection``, since the property may be this
         # current batch.
         response = self._client._base_connection._make_request(
-            "POST", url, data=body, headers=headers
+            "POST", url, data=body, headers=headers, timeout=timeout
         )
         responses = list(_unpack_batch_response(response))
         self._finish_futures(responses)
@@ -313,7 +322,7 @@ def _unpack_batch_response(response):
     parser = Parser()
     message = _generate_faux_mime_message(parser, response)
 
-    if not isinstance(message._payload, list):
+    if not isinstance(message._payload, list):  # pragma: NO COVER
         raise ValueError("Bad response:  not multi-part")
 
     for subrequest in message._payload:

storage/setup.py

@@ -29,8 +29,8 @@
 # 'Development Status :: 5 - Production/Stable'
 release_status = "Development Status :: 5 - Production/Stable"
 dependencies = [
-    "google-auth >= 1.2.0",
-    "google-cloud-core >= 1.0.3, < 2.0dev",
+    "google-auth >= 1.9.0, < 2.0dev",
+    "google-cloud-core >= 1.1.0, < 2.0dev",
     "google-resumable-media >= 0.5.0, < 0.6dev",
 ]
 extras = {}

storage/tests/unit/test__http.py

@@ -51,7 +51,11 @@ def test_extra_headers(self):
         }
         expected_uri = conn.build_api_url("/rainbow")
         http.request.assert_called_once_with(
-            data=req_data, headers=expected_headers, method="GET", url=expected_uri
+            data=req_data,
+            headers=expected_headers,
+            method="GET",
+            url=expected_uri,
+            timeout=None,
         )
 
     def test_build_api_url_no_extra_query_params(self):

storage/tests/unit/test_batch.py

@@ -147,7 +147,7 @@ def test__make_request_GET_normal(self):
         # Check the queued request
         self.assertEqual(len(batch._requests), 1)
         request = batch._requests[0]
-        request_method, request_url, _, request_data = request
+        request_method, request_url, _, request_data, _ = request
         self.assertEqual(request_method, "GET")
         self.assertEqual(request_url, url)
         self.assertIsNone(request_data)
@@ -174,7 +174,7 @@ def test__make_request_POST_normal(self):
         http.request.assert_not_called()
 
         request = batch._requests[0]
-        request_method, request_url, _, request_data = request
+        request_method, request_url, _, request_data, _ = request
         self.assertEqual(request_method, "POST")
         self.assertEqual(request_url, url)
         self.assertEqual(request_data, data)
@@ -201,7 +201,7 @@ def test__make_request_PATCH_normal(self):
         http.request.assert_not_called()
 
         request = batch._requests[0]
-        request_method, request_url, _, request_data = request
+        request_method, request_url, _, request_data, _ = request
         self.assertEqual(request_method, "PATCH")
         self.assertEqual(request_url, url)
         self.assertEqual(request_data, data)
@@ -228,7 +228,7 @@ def test__make_request_DELETE_normal(self):
         # Check the queued request
         self.assertEqual(len(batch._requests), 1)
         request = batch._requests[0]
-        request_method, request_url, _, request_data = request
+        request_method, request_url, _, request_data, _ = request
         self.assertEqual(request_method, "DELETE")
         self.assertEqual(request_url, url)
         self.assertIsNone(request_data)
@@ -340,7 +340,11 @@ def test_finish_nonempty(self):
 
         expected_url = "{}/batch/storage/v1".format(batch.API_BASE_URL)
         http.request.assert_called_once_with(
-            method="POST", url=expected_url, headers=mock.ANY, data=mock.ANY
+            method="POST",
+            url=expected_url,
+            headers=mock.ANY,
+            data=mock.ANY,
+            timeout=mock.ANY,
         )
 
         request_info = self._get_mutlipart_request(http)
@@ -406,7 +410,11 @@ def test_finish_nonempty_with_status_failure(self):
 
         expected_url = "{}/batch/storage/v1".format(batch.API_BASE_URL)
         http.request.assert_called_once_with(
-            method="POST", url=expected_url, headers=mock.ANY, data=mock.ANY
+            method="POST",
+            url=expected_url,
+            headers=mock.ANY,
+            data=mock.ANY,
+            timeout=mock.ANY,
         )
 
         _, request_body, _, boundary = self._get_mutlipart_request(http)
@@ -620,8 +628,10 @@ class _Connection(object):
     def __init__(self, **kw):
         self.__dict__.update(kw)
 
-    def _make_request(self, method, url, data=None, headers=None):
-        return self.http.request(url=url, method=method, headers=headers, data=data)
+    def _make_request(self, method, url, data=None, headers=None, timeout=None):
+        return self.http.request(
+            url=url, method=method, headers=headers, data=data, timeout=timeout
+        )
 
 
 class _MockObject(object):

storage/tests/unit/test_client.py

@@ -271,7 +271,7 @@ def test_get_service_account_email_wo_project(self):
             ]
         )
         http.request.assert_called_once_with(
-            method="GET", url=URI, data=None, headers=mock.ANY
+            method="GET", url=URI, data=None, headers=mock.ANY, timeout=mock.ANY
         )
 
     def test_get_service_account_email_w_project(self):
@@ -297,7 +297,7 @@ def test_get_service_account_email_w_project(self):
             ]
         )
         http.request.assert_called_once_with(
-            method="GET", url=URI, data=None, headers=mock.ANY
+            method="GET", url=URI, data=None, headers=mock.ANY, timeout=mock.ANY
         )
 
     def test_bucket(self):
@@ -366,7 +366,7 @@ def test_get_bucket_with_string_miss(self):
             client.get_bucket(NONESUCH)
 
         http.request.assert_called_once_with(
-            method="GET", url=URI, data=mock.ANY, headers=mock.ANY
+            method="GET", url=URI, data=mock.ANY, headers=mock.ANY, timeout=mock.ANY
         )
 
     def test_get_bucket_with_string_hit(self):
@@ -396,7 +396,7 @@ def test_get_bucket_with_string_hit(self):
         self.assertIsInstance(bucket, Bucket)
         self.assertEqual(bucket.name, BUCKET_NAME)
         http.request.assert_called_once_with(
-            method="GET", url=URI, data=mock.ANY, headers=mock.ANY
+            method="GET", url=URI, data=mock.ANY, headers=mock.ANY, timeout=mock.ANY
         )
 
     def test_get_bucket_with_object_miss(self):
@@ -427,7 +427,7 @@ def test_get_bucket_with_object_miss(self):
             client.get_bucket(bucket_obj)
 
         http.request.assert_called_once_with(
-            method="GET", url=URI, data=mock.ANY, headers=mock.ANY
+            method="GET", url=URI, data=mock.ANY, headers=mock.ANY, timeout=mock.ANY
         )
 
     def test_get_bucket_with_object_hit(self):
@@ -458,7 +458,7 @@ def test_get_bucket_with_object_hit(self):
         self.assertIsInstance(bucket, Bucket)
         self.assertEqual(bucket.name, bucket_name)
         http.request.assert_called_once_with(
-            method="GET", url=URI, data=mock.ANY, headers=mock.ANY
+            method="GET", url=URI, data=mock.ANY, headers=mock.ANY, timeout=mock.ANY
         )
 
     def test_lookup_bucket_miss(self):
@@ -485,7 +485,7 @@ def test_lookup_bucket_miss(self):
 
         self.assertIsNone(bucket)
         http.request.assert_called_once_with(
-            method="GET", url=URI, data=mock.ANY, headers=mock.ANY
+            method="GET", url=URI, data=mock.ANY, headers=mock.ANY, timeout=mock.ANY
         )
 
     def test_lookup_bucket_hit(self):
@@ -514,7 +514,7 @@ def test_lookup_bucket_hit(self):
         self.assertIsInstance(bucket, Bucket)
         self.assertEqual(bucket.name, BUCKET_NAME)
         http.request.assert_called_once_with(
-            method="GET", url=URI, data=mock.ANY, headers=mock.ANY
+            method="GET", url=URI, data=mock.ANY, headers=mock.ANY, timeout=mock.ANY
         )
 
     def test_create_bucket_w_missing_client_project(self):
@@ -666,7 +666,7 @@ def test_create_bucket_w_string_success(self):
         self.assertEqual(bucket.name, bucket_name)
         self.assertTrue(bucket.requester_pays)
         http.request.assert_called_once_with(
-            method="POST", url=URI, data=mock.ANY, headers=mock.ANY
+            method="POST", url=URI, data=mock.ANY, headers=mock.ANY, timeout=mock.ANY
         )
         json_sent = http.request.call_args_list[0][1]["data"]
         self.assertEqual(json_expected, json.loads(json_sent))
@@ -706,7 +706,7 @@ def test_create_bucket_w_object_success(self):
         self.assertEqual(bucket.name, bucket_name)
         self.assertTrue(bucket.requester_pays)
         http.request.assert_called_once_with(
-            method="POST", url=URI, data=mock.ANY, headers=mock.ANY
+            method="POST", url=URI, data=mock.ANY, headers=mock.ANY, timeout=mock.ANY
         )
         json_sent = http.request.call_args_list[0][1]["data"]
         self.assertEqual(json_expected, json.loads(json_sent))
@@ -848,7 +848,11 @@ def test_list_buckets_empty(self):
         self.assertEqual(len(buckets), 0)
 
         http.request.assert_called_once_with(
-            method="GET", url=mock.ANY, data=mock.ANY, headers=mock.ANY
+            method="GET",
+            url=mock.ANY,
+            data=mock.ANY,
+            headers=mock.ANY,
+            timeout=mock.ANY,
         )
 
         requested_url = http.request.mock_calls[0][2]["url"]
@@ -883,7 +887,11 @@ def test_list_buckets_explicit_project(self):
         self.assertEqual(len(buckets), 0)
 
         http.request.assert_called_once_with(
-            method="GET", url=mock.ANY, data=mock.ANY, headers=mock.ANY
+            method="GET",
+            url=mock.ANY,
+            data=mock.ANY,
+            headers=mock.ANY,
+            timeout=mock.ANY,
         )
 
         requested_url = http.request.mock_calls[0][2]["url"]
@@ -918,7 +926,11 @@ def test_list_buckets_non_empty(self):
         self.assertEqual(buckets[0].name, BUCKET_NAME)
 
         http.request.assert_called_once_with(
-            method="GET", url=mock.ANY, data=mock.ANY, headers=mock.ANY
+            method="GET",
+            url=mock.ANY,
+            data=mock.ANY,
+            headers=mock.ANY,
+            timeout=mock.ANY,
         )
 
     def test_list_buckets_all_arguments(self):
@@ -948,7 +960,11 @@ def test_list_buckets_all_arguments(self):
         buckets = list(iterator)
         self.assertEqual(buckets, [])
         http.request.assert_called_once_with(
-            method="GET", url=mock.ANY, data=mock.ANY, headers=mock.ANY
+            method="GET",
+            url=mock.ANY,
+            data=mock.ANY,
+            headers=mock.ANY,
+            timeout=mock.ANY,
         )
 
         requested_url = http.request.mock_calls[0][2]["url"]
@@ -1077,7 +1093,7 @@ def _create_hmac_key_helper(self, explicit_project=None, user_project=None):
 
         FULL_URI = "{}?{}".format(URI, urlencode(qs_params))
         http.request.assert_called_once_with(
-            method="POST", url=FULL_URI, data=None, headers=mock.ANY
+            method="POST", url=FULL_URI, data=None, headers=mock.ANY, timeout=mock.ANY
         )
 
     def test_create_hmac_key_defaults(self):
@@ -1112,7 +1128,7 @@ def test_list_hmac_keys_defaults_empty(self):
             ]
         )
         http.request.assert_called_once_with(
-            method="GET", url=URI, data=None, headers=mock.ANY
+            method="GET", url=URI, data=None, headers=mock.ANY, timeout=mock.ANY
         )
 
     def test_list_hmac_keys_explicit_non_empty(self):
@@ -1176,7 +1192,7 @@ def test_list_hmac_keys_explicit_non_empty(self):
             "userProject": USER_PROJECT,
         }
         http.request.assert_called_once_with(
-            method="GET", url=mock.ANY, data=None, headers=mock.ANY
+            method="GET", url=mock.ANY, data=None, headers=mock.ANY, timeout=mock.ANY
         )
         kwargs = http.request.mock_calls[0].kwargs
         uri = kwargs["url"]
@@ -1223,7 +1239,7 @@ def test_get_hmac_key_metadata_wo_project(self):
             ]
         )
         http.request.assert_called_once_with(
-            method="GET", url=URI, data=None, headers=mock.ANY
+            method="GET", url=URI, data=None, headers=mock.ANY, timeout=mock.ANY
         )
 
     def test_get_hmac_key_metadata_w_project(self):
@@ -1273,5 +1289,5 @@ def test_get_hmac_key_metadata_w_project(self):
         FULL_URI = "{}?{}".format(URI, urlencode(qs_params))
 
         http.request.assert_called_once_with(
-            method="GET", url=FULL_URI, data=None, headers=mock.ANY
+            method="GET", url=FULL_URI, data=None, headers=mock.ANY, timeout=mock.ANY
         )