Skip to content

Commit d41fd5c

Browse files
tswasttswast
authored
chore: execute retry tests serially, since they depend on global time (#2265)
* chore: migrate test_base retry tests * migrate job_helpers test * migrate more tests * fix initiate resumable upload tests * fix failing tests * remove dead test code
1 parent b684832 commit d41fd5c

12 files changed

+1226
-1001
lines changed

‎tests/unit/conftest.py

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,7 @@
1313
# limitations under the License.
1414

1515
from unittest import mock
16+
import threading
1617

1718
import pytest
1819

@@ -24,6 +25,18 @@ def client():
2425
yield make_client()
2526

2627

28+
time_lock = threading.Lock()
29+
30+
31+
@pytest.fixture
32+
def global_time_lock():
33+
"""Fixture to run tests serially that depend on the global time state,
34+
such as tests of retry behavior.
35+
"""
36+
with time_lock:
37+
yield
38+
39+
2740
@pytest.fixture
2841
def PROJECT():
2942
yield "PROJECT"

‎tests/unit/job/test_async_job_retry.py

Lines changed: 139 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,139 @@
1+
# Copyright 2025 Google LLC
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
15+
from unittest import mock
16+
17+
import google.api_core.retry
18+
from google.api_core import exceptions
19+
20+
from . import helpers
21+
import google.cloud.bigquery.job
22+
23+
24+
PROJECT = "test-project"
25+
JOB_ID = "test-job-id"
26+
27+
28+
def test_cancel_w_custom_retry(global_time_lock):
29+
from google.cloud.bigquery.retry import DEFAULT_RETRY
30+
31+
api_path = "/projects/{}/jobs/{}/cancel".format(PROJECT, JOB_ID)
32+
resource = {
33+
"jobReference": {
34+
"jobId": JOB_ID,
35+
"projectId": PROJECT,
36+
"location": None,
37+
},
38+
"configuration": {"test": True},
39+
}
40+
expected = resource.copy()
41+
expected["statistics"] = {}
42+
response = {"job": resource}
43+
conn = helpers.make_connection(
44+
ValueError,
45+
response,
46+
)
47+
client = helpers._make_client(project=PROJECT, connection=conn)
48+
job = google.cloud.bigquery.job._AsyncJob(
49+
google.cloud.bigquery.job._JobReference(JOB_ID, PROJECT, "EU"), client
50+
)
51+
52+
retry = DEFAULT_RETRY.with_deadline(1).with_predicate(
53+
lambda exc: isinstance(exc, ValueError)
54+
)
55+
56+
with mock.patch(
57+
"google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes"
58+
) as final_attributes:
59+
result = job.cancel(retry=retry, timeout=7.5)
60+
61+
final_attributes.assert_called()
62+
63+
assert result is True
64+
assert job._properties == expected
65+
conn.api_request.assert_has_calls(
66+
[
67+
mock.call(
68+
method="POST",
69+
path=api_path,
70+
query_params={"location": "EU"},
71+
timeout=7.5,
72+
),
73+
mock.call(
74+
method="POST",
75+
path=api_path,
76+
query_params={"location": "EU"},
77+
timeout=7.5,
78+
), # was retried once
79+
],
80+
)
81+
82+
83+
def test_result_w_retry_wo_state(global_time_lock):
84+
from google.cloud.bigquery.retry import DEFAULT_GET_JOB_TIMEOUT
85+
86+
begun_job_resource = helpers._make_job_resource(
87+
job_id=JOB_ID, project_id=PROJECT, location="EU", started=True
88+
)
89+
done_job_resource = helpers._make_job_resource(
90+
job_id=JOB_ID,
91+
project_id=PROJECT,
92+
location="EU",
93+
started=True,
94+
ended=True,
95+
)
96+
conn = helpers.make_connection(
97+
exceptions.NotFound("not normally retriable"),
98+
begun_job_resource,
99+
exceptions.NotFound("not normally retriable"),
100+
done_job_resource,
101+
)
102+
client = helpers._make_client(project=PROJECT, connection=conn)
103+
job = google.cloud.bigquery.job._AsyncJob(
104+
google.cloud.bigquery.job._JobReference(JOB_ID, PROJECT, "EU"), client
105+
)
106+
custom_predicate = mock.Mock()
107+
custom_predicate.return_value = True
108+
custom_retry = google.api_core.retry.Retry(
109+
predicate=custom_predicate,
110+
initial=0.001,
111+
maximum=0.001,
112+
deadline=0.1,
113+
)
114+
assert job.result(retry=custom_retry) is job
115+
116+
begin_call = mock.call(
117+
method="POST",
118+
path=f"/projects/{PROJECT}/jobs",
119+
data={
120+
"jobReference": {
121+
"jobId": JOB_ID,
122+
"projectId": PROJECT,
123+
"location": "EU",
124+
}
125+
},
126+
timeout=None,
127+
)
128+
reload_call = mock.call(
129+
method="GET",
130+
path=f"/projects/{PROJECT}/jobs/{JOB_ID}",
131+
query_params={
132+
"projection": "full",
133+
"location": "EU",
134+
},
135+
timeout=DEFAULT_GET_JOB_TIMEOUT,
136+
)
137+
conn.api_request.assert_has_calls(
138+
[begin_call, begin_call, reload_call, reload_call]
139+
)

‎tests/unit/job/test_base.py

Lines changed: 0 additions & 104 deletions
Original file line numberDiff line numberDiff line change
@@ -17,8 +17,6 @@
1717
import unittest
1818
from unittest import mock
1919

20-
from google.api_core import exceptions
21-
import google.api_core.retry
2220
from google.api_core.future import polling
2321
import pytest
2422

@@ -882,50 +880,6 @@ def test_cancel_explicit(self):
882880
)
883881
self.assertEqual(job._properties, expected)
884882

885-
def test_cancel_w_custom_retry(self):
886-
from google.cloud.bigquery.retry import DEFAULT_RETRY
887-
888-
api_path = "/projects/{}/jobs/{}/cancel".format(self.PROJECT, self.JOB_ID)
889-
resource = {
890-
"jobReference": {
891-
"jobId": self.JOB_ID,
892-
"projectId": self.PROJECT,
893-
"location": None,
894-
},
895-
"configuration": {"test": True},
896-
}
897-
expected = resource.copy()
898-
expected["statistics"] = {}
899-
response = {"job": resource}
900-
job = self._set_properties_job()
901-
902-
api_request_patcher = mock.patch.object(
903-
job._client._connection, "api_request", side_effect=[ValueError, response]
904-
)
905-
retry = DEFAULT_RETRY.with_deadline(1).with_predicate(
906-
lambda exc: isinstance(exc, ValueError)
907-
)
908-
909-
with api_request_patcher as fake_api_request:
910-
with mock.patch(
911-
"google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes"
912-
) as final_attributes:
913-
result = job.cancel(retry=retry, timeout=7.5)
914-
915-
final_attributes.assert_called()
916-
917-
self.assertTrue(result)
918-
self.assertEqual(job._properties, expected)
919-
self.assertEqual(
920-
fake_api_request.call_args_list,
921-
[
922-
mock.call(method="POST", path=api_path, query_params={}, timeout=7.5),
923-
mock.call(
924-
method="POST", path=api_path, query_params={}, timeout=7.5
925-
), # was retried once
926-
],
927-
)
928-
929883
def test__set_future_result_wo_done(self):
930884
client = _make_client(project=self.PROJECT)
931885
job = self._make_one(self.JOB_ID, client)
@@ -1069,64 +1023,6 @@ def test_result_default_wo_state(self):
10691023
)
10701024
conn.api_request.assert_has_calls([begin_call, begin_call, reload_call])
10711025

1072-
def test_result_w_retry_wo_state(self):
1073-
from google.cloud.bigquery.retry import DEFAULT_GET_JOB_TIMEOUT
1074-
1075-
begun_job_resource = _make_job_resource(
1076-
job_id=self.JOB_ID, project_id=self.PROJECT, location="EU", started=True
1077-
)
1078-
done_job_resource = _make_job_resource(
1079-
job_id=self.JOB_ID,
1080-
project_id=self.PROJECT,
1081-
location="EU",
1082-
started=True,
1083-
ended=True,
1084-
)
1085-
conn = make_connection(
1086-
exceptions.NotFound("not normally retriable"),
1087-
begun_job_resource,
1088-
exceptions.NotFound("not normally retriable"),
1089-
done_job_resource,
1090-
)
1091-
client = _make_client(project=self.PROJECT, connection=conn)
1092-
job = self._make_one(
1093-
self._job_reference(self.JOB_ID, self.PROJECT, "EU"), client
1094-
)
1095-
custom_predicate = mock.Mock()
1096-
custom_predicate.return_value = True
1097-
custom_retry = google.api_core.retry.Retry(
1098-
predicate=custom_predicate,
1099-
initial=0.001,
1100-
maximum=0.001,
1101-
deadline=0.1,
1102-
)
1103-
self.assertIs(job.result(retry=custom_retry), job)
1104-
1105-
begin_call = mock.call(
1106-
method="POST",
1107-
path=f"/projects/{self.PROJECT}/jobs",
1108-
data={
1109-
"jobReference": {
1110-
"jobId": self.JOB_ID,
1111-
"projectId": self.PROJECT,
1112-
"location": "EU",
1113-
}
1114-
},
1115-
timeout=None,
1116-
)
1117-
reload_call = mock.call(
1118-
method="GET",
1119-
path=f"/projects/{self.PROJECT}/jobs/{self.JOB_ID}",
1120-
query_params={
1121-
"projection": "full",
1122-
"location": "EU",
1123-
},
1124-
timeout=DEFAULT_GET_JOB_TIMEOUT,
1125-
)
1126-
conn.api_request.assert_has_calls(
1127-
[begin_call, begin_call, reload_call, reload_call]
1128-
)
1129-
11301026
def test_result_explicit_w_state(self):
11311027
conn = make_connection()
11321028
client = _make_client(project=self.PROJECT, connection=conn)

0 commit comments

Comments
 (0)