chore: migrate test_base retry tests

Author: tswast

tests/unit/conftest.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 
 from unittest import mock
+import threading
 
 import pytest
 
@@ -24,6 +25,18 @@ def client():
     yield make_client()
 
 
+time_lock = threading.Lock()
+
+
+@pytest.fixture
+def global_time_lock():
+    """Fixture to run tests serially that depend on the global time state,
+    such as tests of retry behavior.
+    """
+    with time_lock:
+        yield
+
+
 @pytest.fixture
 def PROJECT():
     yield "PROJECT"

tests/unit/job/test_async_job_retry.py

@@ -0,0 +1,139 @@
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from unittest import mock
+
+import google.api_core.retry
+from google.api_core import exceptions
+
+from . import helpers
+import google.cloud.bigquery.job
+
+
+PROJECT = "test-project"
+JOB_ID = "test-job-id"
+
+
+def test_cancel_w_custom_retry(global_time_lock):
+    from google.cloud.bigquery.retry import DEFAULT_RETRY
+
+    api_path = "/projects/{}/jobs/{}/cancel".format(PROJECT, JOB_ID)
+    resource = {
+        "jobReference": {
+            "jobId": JOB_ID,
+            "projectId": PROJECT,
+            "location": None,
+        },
+        "configuration": {"test": True},
+    }
+    expected = resource.copy()
+    expected["statistics"] = {}
+    response = {"job": resource}
+    conn = helpers.make_connection(
+        ValueError,
+        response,
+    )
+    client = helpers._make_client(project=PROJECT, connection=conn)
+    job = google.cloud.bigquery.job._AsyncJob(
+        google.cloud.bigquery.job._JobReference(JOB_ID, PROJECT, "EU"), client
+    )
+
+    retry = DEFAULT_RETRY.with_deadline(1).with_predicate(
+        lambda exc: isinstance(exc, ValueError)
+    )
+
+    with mock.patch(
+        "google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes"
+    ) as final_attributes:
+        result = job.cancel(retry=retry, timeout=7.5)
+
+    final_attributes.assert_called()
+
+    assert result is True
+    assert job._properties == expected
+    conn.api_request.assert_has_calls(
+        [
+            mock.call(
+                method="POST",
+                path=api_path,
+                query_params={"location": "EU"},
+                timeout=7.5,
+            ),
+            mock.call(
+                method="POST",
+                path=api_path,
+                query_params={"location": "EU"},
+                timeout=7.5,
+            ),  # was retried once
+        ],
+    )
+
+
+def test_result_w_retry_wo_state(global_time_lock):
+    from google.cloud.bigquery.retry import DEFAULT_GET_JOB_TIMEOUT
+
+    begun_job_resource = helpers._make_job_resource(
+        job_id=JOB_ID, project_id=PROJECT, location="EU", started=True
+    )
+    done_job_resource = helpers._make_job_resource(
+        job_id=JOB_ID,
+        project_id=PROJECT,
+        location="EU",
+        started=True,
+        ended=True,
+    )
+    conn = helpers.make_connection(
+        exceptions.NotFound("not normally retriable"),
+        begun_job_resource,
+        exceptions.NotFound("not normally retriable"),
+        done_job_resource,
+    )
+    client = helpers._make_client(project=PROJECT, connection=conn)
+    job = google.cloud.bigquery.job._AsyncJob(
+        google.cloud.bigquery.job._JobReference(JOB_ID, PROJECT, "EU"), client
+    )
+    custom_predicate = mock.Mock()
+    custom_predicate.return_value = True
+    custom_retry = google.api_core.retry.Retry(
+        predicate=custom_predicate,
+        initial=0.001,
+        maximum=0.001,
+        deadline=0.1,
+    )
+    assert job.result(retry=custom_retry) is job
+
+    begin_call = mock.call(
+        method="POST",
+        path=f"/projects/{PROJECT}/jobs",
+        data={
+            "jobReference": {
+                "jobId": JOB_ID,
+                "projectId": PROJECT,
+                "location": "EU",
+            }
+        },
+        timeout=None,
+    )
+    reload_call = mock.call(
+        method="GET",
+        path=f"/projects/{PROJECT}/jobs/{JOB_ID}",
+        query_params={
+            "projection": "full",
+            "location": "EU",
+        },
+        timeout=DEFAULT_GET_JOB_TIMEOUT,
+    )
+    conn.api_request.assert_has_calls(
+        [begin_call, begin_call, reload_call, reload_call]
+    )

tests/unit/job/test_base.py

@@ -17,8 +17,6 @@
 import unittest
 from unittest import mock
 
-from google.api_core import exceptions
-import google.api_core.retry
 from google.api_core.future import polling
 import pytest
 
@@ -882,50 +880,6 @@ def test_cancel_explicit(self):
         )
         self.assertEqual(job._properties, expected)
 
-    def test_cancel_w_custom_retry(self):
-        from google.cloud.bigquery.retry import DEFAULT_RETRY
-
-        api_path = "/projects/{}/jobs/{}/cancel".format(self.PROJECT, self.JOB_ID)
-        resource = {
-            "jobReference": {
-                "jobId": self.JOB_ID,
-                "projectId": self.PROJECT,
-                "location": None,
-            },
-            "configuration": {"test": True},
-        }
-        expected = resource.copy()
-        expected["statistics"] = {}
-        response = {"job": resource}
-        job = self._set_properties_job()
-
-        api_request_patcher = mock.patch.object(
-            job._client._connection, "api_request", side_effect=[ValueError, response]
-        )
-        retry = DEFAULT_RETRY.with_deadline(1).with_predicate(
-            lambda exc: isinstance(exc, ValueError)
-        )
-
-        with api_request_patcher as fake_api_request:
-            with mock.patch(
-                "google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes"
-            ) as final_attributes:
-                result = job.cancel(retry=retry, timeout=7.5)
-
-            final_attributes.assert_called()
-
-        self.assertTrue(result)
-        self.assertEqual(job._properties, expected)
-        self.assertEqual(
-            fake_api_request.call_args_list,
-            [
-                mock.call(method="POST", path=api_path, query_params={}, timeout=7.5),
-                mock.call(
-                    method="POST", path=api_path, query_params={}, timeout=7.5
-                ),  # was retried once
-            ],
-        )
-
     def test__set_future_result_wo_done(self):
         client = _make_client(project=self.PROJECT)
         job = self._make_one(self.JOB_ID, client)
@@ -1069,64 +1023,6 @@ def test_result_default_wo_state(self):
         )
         conn.api_request.assert_has_calls([begin_call, begin_call, reload_call])
 
-    def test_result_w_retry_wo_state(self):
-        from google.cloud.bigquery.retry import DEFAULT_GET_JOB_TIMEOUT
-
-        begun_job_resource = _make_job_resource(
-            job_id=self.JOB_ID, project_id=self.PROJECT, location="EU", started=True
-        )
-        done_job_resource = _make_job_resource(
-            job_id=self.JOB_ID,
-            project_id=self.PROJECT,
-            location="EU",
-            started=True,
-            ended=True,
-        )
-        conn = make_connection(
-            exceptions.NotFound("not normally retriable"),
-            begun_job_resource,
-            exceptions.NotFound("not normally retriable"),
-            done_job_resource,
-        )
-        client = _make_client(project=self.PROJECT, connection=conn)
-        job = self._make_one(
-            self._job_reference(self.JOB_ID, self.PROJECT, "EU"), client
-        )
-        custom_predicate = mock.Mock()
-        custom_predicate.return_value = True
-        custom_retry = google.api_core.retry.Retry(
-            predicate=custom_predicate,
-            initial=0.001,
-            maximum=0.001,
-            deadline=0.1,
-        )
-        self.assertIs(job.result(retry=custom_retry), job)
-
-        begin_call = mock.call(
-            method="POST",
-            path=f"/projects/{self.PROJECT}/jobs",
-            data={
-                "jobReference": {
-                    "jobId": self.JOB_ID,
-                    "projectId": self.PROJECT,
-                    "location": "EU",
-                }
-            },
-            timeout=None,
-        )
-        reload_call = mock.call(
-            method="GET",
-            path=f"/projects/{self.PROJECT}/jobs/{self.JOB_ID}",
-            query_params={
-                "projection": "full",
-                "location": "EU",
-            },
-            timeout=DEFAULT_GET_JOB_TIMEOUT,
-        )
-        conn.api_request.assert_has_calls(
-            [begin_call, begin_call, reload_call, reload_call]
-        )
-
     def test_result_explicit_w_state(self):
         conn = make_connection()
         client = _make_client(project=self.PROJECT, connection=conn)

tests/unit/job/test_query.py

@@ -1335,7 +1335,7 @@ def test_result_with_max_results(self):
             [jobs_get_call, query_page_waiting_call, query_page_2_call]
         )
 
-    def test_result_w_custom_retry(self):
+    def test_result_w_custom_retry(self, global_time_lock):
         from google.cloud.bigquery.table import RowIterator
 
         query_resource = {
@@ -1455,7 +1455,7 @@ def test_result_w_empty_schema(self):
         self.assertEqual(result.location, "asia-northeast1")
         self.assertEqual(result.query_id, "xyz-abc")
 
-    def test_result_w_timeout_doesnt_raise(self):
+    def test_result_w_timeout_doesnt_raise(self, global_time_lock):
         import google.cloud.bigquery.client
 
         begun_resource = self._make_resource()
@@ -1505,7 +1505,7 @@ def test_result_w_timeout_doesnt_raise(self):
             ]
         )
 
-    def test_result_w_timeout_raises_concurrent_futures_timeout(self):
+    def test_result_w_timeout_raises_concurrent_futures_timeout(self, global_time_lock):
         import google.cloud.bigquery.client
 
         begun_resource = self._make_resource()

tests/unit/test__job_helpers.py

@@ -450,7 +450,7 @@ def test_query_and_wait_uses_jobs_insert():
     )
 
 
-def test_query_and_wait_retries_job():
+def test_query_and_wait_retries_job(global_time_lock):
     freezegun.freeze_time(auto_tick_seconds=100)
     client = mock.create_autospec(Client)
     client._call_api.__name__ = "_call_api"
@@ -512,8 +512,8 @@ def test_query_and_wait_retries_job():
         assert kwargs["path"] == request_path
 
 
-@freezegun.freeze_time(auto_tick_seconds=100)
-def test_query_and_wait_retries_job_times_out():
+def test_query_and_wait_retries_job_times_out(global_time_lock):
+    freezegun.freeze_time(auto_tick_seconds=100)
     client = mock.create_autospec(Client)
     client._call_api.__name__ = "_call_api"
     client._call_api.__qualname__ = "Client._call_api"

tests/unit/test_client.py

@@ -313,7 +313,7 @@ def test__call_api_extra_headers(self):
         headers = kwargs["headers"]
         assert headers["x-goog-request-reason"] == "because-friday"
 
-    def test__call_api_applying_custom_retry_on_timeout(self):
+    def test__call_api_applying_custom_retry_on_timeout(self, global_time_lock):
         from concurrent.futures import TimeoutError
         from google.cloud.bigquery.retry import DEFAULT_RETRY
 
@@ -644,7 +644,7 @@ def test_get_service_account_email_w_alternate_project(self):
         )
         self.assertEqual(service_account_email, email)
 
-    def test_get_service_account_email_w_custom_retry(self):
+    def test_get_service_account_email_w_custom_retry(self, global_time_lock):
         from google.cloud.bigquery.retry import DEFAULT_RETRY
 
         api_path = "/projects/{}/serviceAccount".format(self.PROJECT)
@@ -3941,7 +3941,7 @@ def test__initiate_resumable_upload(self):
     def test__initiate_resumable_upload_mtls(self):
         self._initiate_resumable_upload_helper(mtls=True)
 
-    def test__initiate_resumable_upload_with_retry(self):
+    def test__initiate_resumable_upload_with_retry(self, global_time_lock):
         self._initiate_resumable_upload_helper(num_retries=11)
 
     def _do_multipart_upload_success_helper(
@@ -4011,7 +4011,7 @@ def test__do_multipart_upload_mtls(self, get_boundary):
         self._do_multipart_upload_success_helper(get_boundary, mtls=True)
 
     @mock.patch("google.resumable_media._upload.get_boundary", return_value=b"==0==")
-    def test__do_multipart_upload_with_retry(self, get_boundary):
+    def test__do_multipart_upload_with_retry(self, get_boundary, global_time_lock):
         self._do_multipart_upload_success_helper(get_boundary, num_retries=8)
 
     @mock.patch("google.resumable_media._upload.get_boundary", return_value=b"==0==")
@@ -5602,7 +5602,7 @@ def test_query_job_rpc_fail_w_conflict_random_id_job_fetch_retries_404(self):
         assert result.job_id == job_id
 
     def test_query_job_rpc_fail_w_conflict_random_id_job_fetch_retries_404_and_query_job_insert(
-        self,
+        self, global_time_lock
     ):
         """Regression test for https://github.com/googleapis/python-bigquery/issues/2134
 

tests/unit/test_client_bigframes.py

@@ -338,7 +338,7 @@ def test_query_and_wait_bigframes_with_jobs_insert_dry_run_no_callback(client):
     assert result.schema == [bigquery.SchemaField("_f0", "INTEGER")]
 
 
-def test_query_and_wait_bigframes_with_query_retry_callbacks(client):
+def test_query_and_wait_bigframes_with_query_retry_callbacks(client, global_time_lock):
     created = datetime.datetime(
         2025, 8, 18, 10, 11, 12, 345000, tzinfo=datetime.timezone.utc
     )