Add 'retry' argument to '_AsyncJob.result'. (googleapis#6302)

tseaver · web-flow · commit be78a968ba7c · 2018-11-30T19:27:11.000-05:00
Pass it through to the '_begin' call.  Note that we need to modify
the 'api_core...PollingFuture' class before we can safely pass the
'retry' through to its 'result'.
diff --git a/bigquery/google/cloud/bigquery/job.py b/bigquery/google/cloud/bigquery/job.py
@@ -678,14 +678,17 @@ def done(self, retry=DEFAULT_RETRY):
             self.reload(retry=retry)
         return self.state == _DONE_STATE
 
-    def result(self, timeout=None):
+    def result(self, timeout=None, retry=DEFAULT_RETRY):
         """Start the job and wait for it to complete and get the result.
 
         :type timeout: float
         :param timeout:
             How long (in seconds) to wait for job to complete before raising
             a :class:`concurrent.futures.TimeoutError`.
 
+        :type retry: :class:`google.api_core.retry.Retry`
+        :param retry: (Optional) How to retry the RPC.
+
         :rtype: _AsyncJob
         :returns: This instance.
 
@@ -695,7 +698,7 @@ def result(self, timeout=None):
             not complete in the given timeout.
         """
         if self.state is None:
-            self._begin()
+            self._begin(retry=retry)
         # TODO: modify PollingFuture so it can pass a retry argument to done().
         return super(_AsyncJob, self).result(timeout=timeout)
 
diff --git a/bigquery/tests/unit/test_job.py b/bigquery/tests/unit/test_job.py
@@ -789,13 +789,27 @@ def test_done_already(self):
 
     @mock.patch("google.api_core.future.polling.PollingFuture.result")
     def test_result_default_wo_state(self, result):
+        from google.cloud.bigquery.retry import DEFAULT_RETRY
+
         client = _make_client(project=self.PROJECT)
         job = self._make_one(self.JOB_ID, client)
         begin = job._begin = mock.Mock()
 
         self.assertIs(job.result(), result.return_value)
 
-        begin.assert_called_once()
+        begin.assert_called_once_with(retry=DEFAULT_RETRY)
+        result.assert_called_once_with(timeout=None)
+
+    @mock.patch('google.api_core.future.polling.PollingFuture.result')
+    def test_result_w_retry_wo_state(self, result):
+        client = _make_client(project=self.PROJECT)
+        job = self._make_one(self.JOB_ID, client)
+        begin = job._begin = mock.Mock()
+        retry = mock.Mock()
+
+        self.assertIs(job.result(retry=retry), result.return_value)
+
+        begin.assert_called_once_with(retry=retry)
         result.assert_called_once_with(timeout=None)
 
     @mock.patch("google.api_core.future.polling.PollingFuture.result")
