Add timeout param to JSONConnection.api_request()

Author: plamut

core/google/cloud/_http.py

@@ -222,6 +222,7 @@ def _make_request(
         content_type=None,
         headers=None,
         target_object=None,
+        timeout=None,
     ):
         """A low level method to send a request to the API.
 
@@ -250,6 +251,13 @@ def _make_request(
             custom behavior, for example, to 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: :class:`requests.Response`
         :returns: The HTTP response.
         """
@@ -263,10 +271,12 @@ def _make_request(
         headers[CLIENT_INFO_HEADER] = self.user_agent
         headers["User-Agent"] = self.user_agent
 
-        return self._do_request(method, url, headers, data, target_object)
+        return self._do_request(
+            method, url, headers, data, target_object, timeout=timeout
+        )
 
     def _do_request(
-        self, method, url, headers, data, target_object
+        self, method, url, headers, data, target_object, timeout=None
     ):  # pylint: disable=unused-argument
         """Low-level helper:  perform the actual API request over HTTP.
 
@@ -289,10 +299,19 @@ def _do_request(
             (Optional) Unused ``target_object`` here but may be used by a
             superclass.
 
+        :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: :class:`requests.Response`
         :returns: The HTTP response.
         """
-        return self.http.request(url=url, method=method, headers=headers, data=data)
+        return self.http.request(
+            url=url, method=method, headers=headers, data=data, timeout=timeout
+        )
 
     def api_request(
         self,
@@ -306,6 +325,7 @@ def api_request(
         api_version=None,
         expect_json=True,
         _target_object=None,
+        timeout=None,
     ):
         """Make a request over the HTTP transport to the API.
 
@@ -360,6 +380,13 @@ def api_request(
             can allow custom behavior, for example, to 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.
+
         :raises ~google.cloud.exceptions.GoogleCloudError: if the response code
             is not 200 OK.
         :raises ValueError: if the response content type is not JSON.
@@ -387,6 +414,7 @@ def api_request(
             content_type=content_type,
             headers=headers,
             target_object=_target_object,
+            timeout=timeout,
         )
 
         if not 200 <= response.status_code < 300:

core/tests/unit/test__http.py

@@ -217,7 +217,7 @@ def test__make_request_no_data_no_content_type_no_headers(self):
             CLIENT_INFO_HEADER: conn.user_agent,
         }
         http.request.assert_called_once_with(
-            method="GET", url=url, headers=expected_headers, data=None
+            method="GET", url=url, headers=expected_headers, data=None, timeout=None
         )
 
     def test__make_request_w_data_no_extra_headers(self):
@@ -238,7 +238,7 @@ def test__make_request_w_data_no_extra_headers(self):
             CLIENT_INFO_HEADER: conn.user_agent,
         }
         http.request.assert_called_once_with(
-            method="GET", url=url, headers=expected_headers, data=data
+            method="GET", url=url, headers=expected_headers, data=data, timeout=None
         )
 
     def test__make_request_w_extra_headers(self):
@@ -258,7 +258,30 @@ def test__make_request_w_extra_headers(self):
             CLIENT_INFO_HEADER: conn.user_agent,
         }
         http.request.assert_called_once_with(
-            method="GET", url=url, headers=expected_headers, data=None
+            method="GET", url=url, headers=expected_headers, data=None, timeout=None
+        )
+
+    def test__make_request_w_timeout(self):
+        from google.cloud._http import CLIENT_INFO_HEADER
+
+        http = make_requests_session([make_response()])
+        client = mock.Mock(_http=http, spec=["_http"])
+        conn = self._make_one(client)
+
+        url = "http://example.com/test"
+        conn._make_request("GET", url, timeout=(5.5, 2.8))
+
+        expected_headers = {
+            "Accept-Encoding": "gzip",
+            "User-Agent": conn.user_agent,
+            CLIENT_INFO_HEADER: conn.user_agent,
+        }
+        http.request.assert_called_once_with(
+            method="GET",
+            url=url,
+            headers=expected_headers,
+            data=None,
+            timeout=(5.5, 2.8),
         )
 
     def test_api_request_defaults(self):
@@ -282,7 +305,11 @@ def test_api_request_defaults(self):
             base=conn.API_BASE_URL, version=conn.API_VERSION, path=path
         )
         http.request.assert_called_once_with(
-            method="GET", url=expected_url, headers=expected_headers, data=None
+            method="GET",
+            url=expected_url,
+            headers=expected_headers,
+            data=None,
+            timeout=None,
         )
 
     def test_api_request_w_non_json_response(self):
@@ -321,7 +348,11 @@ def test_api_request_w_query_params(self):
             CLIENT_INFO_HEADER: conn.user_agent,
         }
         http.request.assert_called_once_with(
-            method="GET", url=mock.ANY, headers=expected_headers, data=None
+            method="GET",
+            url=mock.ANY,
+            headers=expected_headers,
+            data=None,
+            timeout=None,
         )
 
         url = http.request.call_args[1]["url"]
@@ -351,7 +382,11 @@ def test_api_request_w_headers(self):
             CLIENT_INFO_HEADER: conn.user_agent,
         }
         http.request.assert_called_once_with(
-            method="GET", url=mock.ANY, headers=expected_headers, data=None
+            method="GET",
+            url=mock.ANY,
+            headers=expected_headers,
+            data=None,
+            timeout=None,
         )
 
     def test_api_request_w_extra_headers(self):
@@ -377,7 +412,11 @@ def test_api_request_w_extra_headers(self):
             CLIENT_INFO_HEADER: conn.user_agent,
         }
         http.request.assert_called_once_with(
-            method="GET", url=mock.ANY, headers=expected_headers, data=None
+            method="GET",
+            url=mock.ANY,
+            headers=expected_headers,
+            data=None,
+            timeout=None,
         )
 
     def test_api_request_w_data(self):
@@ -400,7 +439,39 @@ def test_api_request_w_data(self):
         }
 
         http.request.assert_called_once_with(
-            method="POST", url=mock.ANY, headers=expected_headers, data=expected_data
+            method="POST",
+            url=mock.ANY,
+            headers=expected_headers,
+            data=expected_data,
+            timeout=None,
+        )
+
+    def test_api_request_w_timeout(self):
+        from google.cloud._http import CLIENT_INFO_HEADER
+
+        http = make_requests_session(
+            [make_response(content=b"{}", headers=self.JSON_HEADERS)]
+        )
+        client = mock.Mock(_http=http, spec=["_http"])
+        conn = self._make_mock_one(client)
+        path = "/path/required"
+
+        self.assertEqual(conn.api_request("GET", path, timeout=(2.2, 3.3)), {})
+
+        expected_headers = {
+            "Accept-Encoding": "gzip",
+            "User-Agent": conn.user_agent,
+            CLIENT_INFO_HEADER: conn.user_agent,
+        }
+        expected_url = "{base}/mock/{version}{path}".format(
+            base=conn.API_BASE_URL, version=conn.API_VERSION, path=path
+        )
+        http.request.assert_called_once_with(
+            method="GET",
+            url=expected_url,
+            headers=expected_headers,
+            data=None,
+            timeout=(2.2, 3.3),
         )
 
     def test_api_request_w_404(self):