Skip to content

test(spanner): speed up unit test suite and eliminate idle retry sleeps #18582

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
olavloite merged 2 commits into from
Oct 7, 2026
+203 −128
Merged

test(spanner): speed up unit test suite and eliminate idle retry sleeps #18582

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
842ee99
test(spanner): speed up unit test suite and eliminate idle retry sleeps
olavloite Oct 6, 2026
2e25044
chore(spanner): run formatter
olavloite Oct 7, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
30 changes: 23 additions & 7 deletions packages/google-cloud-spanner/tests/unit/_async/test_session.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1219,9 +1219,13 @@ async def unit_of_work(transaction):
pass

await session.create()
await session.run_in_transaction(unit_of_work)
with mock.patch(
"google.cloud.spanner_v1._async._helpers.asyncio.sleep"
) as sleep_mock:
await session.run_in_transaction(unit_of_work)

self.assertEqual(begin_transaction.call_count, 2)
sleep_mock.assert_called_once()

begin_transaction.assert_called_with(
request=BeginTransactionRequest(
Expand Down Expand Up @@ -1261,10 +1265,14 @@ async def unit_of_work(transaction):
pass

await session.create()
await session.run_in_transaction(unit_of_work)
with mock.patch(
"google.cloud.spanner_v1._async._helpers.asyncio.sleep"
) as sleep_mock:
await session.run_in_transaction(unit_of_work)

# Verify retried BeginTransaction API call.
self.assertEqual(begin_transaction.call_count, 2)
sleep_mock.assert_called_once()

begin_transaction.assert_called_with(
request=BeginTransactionRequest(
Expand Down Expand Up @@ -1308,10 +1316,14 @@ async def unit_of_work(transaction):
pass

await session.create()
await session.run_in_transaction(unit_of_work)
with mock.patch(
"google.cloud.spanner_v1._async._helpers.asyncio.sleep"
) as sleep_mock:
await session.run_in_transaction(unit_of_work)

# Verify retried BeginTransaction API call.
self.assertEqual(begin_transaction.call_count, 2)
sleep_mock.assert_called_once()

begin_transaction.assert_called_with(
request=BeginTransactionRequest(
Expand Down Expand Up @@ -1812,7 +1824,7 @@ def _time():
return _results[0]
return 1.0

with mock.patch("time.time", _time):
with mock.patch("google.cloud.spanner_v1._async._helpers.time.time", _time):
with mock.patch(
"google.cloud.spanner_v1._async._helpers.asyncio.sleep",
new_callable=mock.AsyncMock,
Expand Down Expand Up @@ -1893,7 +1905,7 @@ def _time():
return 1.0

with (
mock.patch("time.time", _time),
mock.patch("google.cloud.spanner_v1._async._helpers.time.time", _time),
mock.patch(
"google.cloud.spanner_v1._helpers.random.random", return_value=0
),
Expand Down Expand Up @@ -2817,11 +2829,15 @@ def _time_func():
return 3

# check if current time > deadline
with mock.patch("time.time", _time_func):
with mock.patch(
"google.cloud.spanner_v1._async._helpers.time.time", _time_func
):
with pytest.raises(Exception):
_delay_until_retry(exc_mock, 2, 1, default_retry_delay=0)

with mock.patch("time.time", _time_func):
with mock.patch(
"google.cloud.spanner_v1._async._helpers.time.time", _time_func
):
with mock.patch(
"google.cloud.spanner_v1._helpers._get_retry_delay"
) as get_retry_delay_mock:
Expand Down
7 changes: 6 additions & 1 deletion packages/google-cloud-spanner/tests/unit/_async/test_snapshot.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1367,9 +1367,14 @@ async def test_begin_transaction_retry(self):
TransactionPB(id=TXN_ID),
]

tid = await snapshot._begin_transaction()
with mock.patch(
"google.cloud.spanner_v1._async._helpers.asyncio.sleep"
) as sleep_mock:
tid = await snapshot._begin_transaction()

self.assertEqual(tid, TXN_ID)
self.assertEqual(api.begin_transaction.call_count, 2)
sleep_mock.assert_called_once_with(2)

async def test_update_for_transaction_pb_w_precommit_token(self):
from google.cloud.spanner_v1.types import MultiplexedSessionPrecommitToken
Expand Down
7 changes: 6 additions & 1 deletion packages/google-cloud-spanner/tests/unit/_async/test_transaction_extra.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -135,8 +135,13 @@ async def test_commit_retry_and_precommit_token(self):
final_resp,
]

await txn.commit()
with mock.patch(
"google.cloud.spanner_v1._async._helpers.asyncio.sleep"
) as sleep_mock:
await txn.commit()

self.assertEqual(self.db.spanner_api.commit.call_count, 3)
sleep_mock.assert_called_once_with(2)

async def test_execute_update_request_options_dict(self):
# coverage for line 503
Expand Down
12 changes: 6 additions & 6 deletions packages/google-cloud-spanner/tests/unit/omni/test_login_client.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -41,9 +41,9 @@ def test_login_successful_flow(self):

params = authentication_pb2.HashParameters(
argon2_id_parameters=authentication_pb2.HashParameters.Argon2IdParameters(
iteration_count=3,
memory_usage=64 * 1024,
parallelism=4,
iteration_count=1,
memory_usage=8,
parallelism=1,
hash_size=32,
)
)
Expand Down Expand Up @@ -248,9 +248,9 @@ def test_login_missing_access_token_in_final_response(self):
password = "test_password"
params = authentication_pb2.HashParameters(
argon2_id_parameters=authentication_pb2.HashParameters.Argon2IdParameters(
iteration_count=3,
memory_usage=64 * 1024,
parallelism=4,
iteration_count=1,
memory_usage=8,
parallelism=1,
hash_size=32,
)
)
Expand Down
36 changes: 18 additions & 18 deletions packages/google-cloud-spanner/tests/unit/omni/test_opaque.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -486,9 +486,9 @@ def test_oprf_evaluate(self):
def test_authenticator_validation(self):
valid_params = authentication_pb2.HashParameters(
argon2_id_parameters=authentication_pb2.HashParameters.Argon2IdParameters(
iteration_count=3,
memory_usage=64 * 1024,
parallelism=4,
iteration_count=1,
memory_usage=8,
parallelism=1,
hash_size=32,
)
)
Expand All @@ -512,9 +512,9 @@ def test_authenticator_validation(self):
def test_authenticator_state_errors(self):
params = authentication_pb2.HashParameters(
argon2_id_parameters=authentication_pb2.HashParameters.Argon2IdParameters(
iteration_count=3,
memory_usage=64 * 1024,
parallelism=4,
iteration_count=1,
memory_usage=8,
parallelism=1,
hash_size=32,
)
)
Expand All @@ -541,9 +541,9 @@ def test_authenticator_state_errors(self):
def test_user_authenticator_clear_and_del(self):
params = authentication_pb2.HashParameters(
argon2_id_parameters=authentication_pb2.HashParameters.Argon2IdParameters(
iteration_count=3,
memory_usage=64 * 1024,
parallelism=4,
iteration_count=1,
memory_usage=8,
parallelism=1,
hash_size=32,
)
)
Expand Down Expand Up @@ -584,9 +584,9 @@ def test_full_opaque_handshake_simulation(self):

params = authentication_pb2.HashParameters(
argon2_id_parameters=authentication_pb2.HashParameters.Argon2IdParameters(
iteration_count=3,
memory_usage=64 * 1024,
parallelism=4,
iteration_count=1,
memory_usage=8,
parallelism=1,
hash_size=32,
)
)
Expand Down Expand Up @@ -708,9 +708,9 @@ def test_full_opaque_handshake_simulation(self):
def test_final_request_invalid_masked_response_length(self):
params = authentication_pb2.HashParameters(
argon2_id_parameters=authentication_pb2.HashParameters.Argon2IdParameters(
iteration_count=3,
memory_usage=64 * 1024,
parallelism=4,
iteration_count=1,
memory_usage=8,
parallelism=1,
hash_size=32,
)
)
Expand Down Expand Up @@ -964,9 +964,9 @@ def test_final_request_edge_cases(self):

params = authentication_pb2.HashParameters(
argon2_id_parameters=authentication_pb2.HashParameters.Argon2IdParameters(
iteration_count=3,
memory_usage=64 * 1024,
parallelism=4,
iteration_count=1,
memory_usage=8,
parallelism=1,
hash_size=32,
)
)
Expand Down
17 changes: 15 additions & 2 deletions packages/google-cloud-spanner/tests/unit/test_batch.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -345,12 +345,25 @@ def test_aborted_exception_on_commit_with_retries(self, mock_region):
batch.insert(TABLE_NAME, COLUMNS, VALUES)

# Assertion: Ensure that calling batch.commit() raises Aborted
with self.assertRaises(Aborted) as context:
batch.commit(timeout_secs=1.0, default_retry_delay=0)
delay_call_count = 0

def fake_delay(exc, *args, **kwargs):
nonlocal delay_call_count
delay_call_count += 1
if delay_call_count >= 2:
raise exc

with mock.patch(
"google.cloud.spanner_v1._helpers._delay_until_retry",
side_effect=fake_delay,
):
with self.assertRaises(Aborted) as context:
batch.commit(timeout_secs=1.0, default_retry_delay=0)

# Verify exception includes request_id attribute
self.assertIn("409 Transaction was aborted", str(context.exception))
Comment thread
olavloite marked this conversation as resolved.
self.assertTrue(hasattr(context.exception, "request_id"))
self.assertEqual(delay_call_count, 2)
self.assertGreater(
api.commit.call_count, 1, "commit should be called more than once"
)
Expand Down
39 changes: 24 additions & 15 deletions packages/google-cloud-spanner/tests/unit/test_client.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -59,6 +59,17 @@ def _get_target_class(self):
def _make_one(self, *args, **kwargs):
return self._get_target_class()(*args, **kwargs)

@staticmethod
def _make_instance_admin_api():
from google.cloud.spanner_admin_instance_v1 import InstanceAdminClient
from google.cloud.spanner_admin_instance_v1.services.instance_admin.transports.base import (
InstanceAdminTransport,
)

mock_transport = mock.create_autospec(InstanceAdminTransport, instance=True)
mock_transport._wrapped_methods = {}
return InstanceAdminClient(transport=mock_transport)

def _constructor_test_helper(
self,
expected_scopes,
Expand Down Expand Up @@ -627,15 +638,14 @@ def test_project_name_property(self):

def test_list_instance_configs(self):
from google.cloud.spanner_admin_instance_v1 import (
InstanceAdminClient,
ListInstanceConfigsRequest,
ListInstanceConfigsResponse,
InstanceConfig as InstanceConfigPB,
)
from google.cloud.spanner_admin_instance_v1 import (
InstanceConfig as InstanceConfigPB,
ListInstanceConfigsRequest,
ListInstanceConfigsResponse,
)
Comment thread
olavloite marked this conversation as resolved.

api = InstanceAdminClient(credentials=AnonymousCredentials())
api = self._make_instance_admin_api()
credentials = build_scoped_credentials()
client = self._make_one(project=self.PROJECT, credentials=credentials)
client._instance_admin_api = api
Expand Down Expand Up @@ -676,16 +686,15 @@ def test_list_instance_configs(self):

def test_list_instance_configs_w_options(self):
from google.cloud.spanner_admin_instance_v1 import (
InstanceAdminClient,
ListInstanceConfigsRequest,
ListInstanceConfigsResponse,
InstanceConfig as InstanceConfigPB,
)
from google.cloud.spanner_admin_instance_v1 import (
InstanceConfig as InstanceConfigPB,
ListInstanceConfigsRequest,
ListInstanceConfigsResponse,
)
Comment thread
olavloite marked this conversation as resolved.

credentials = build_scoped_credentials()
api = InstanceAdminClient(credentials=credentials)
api = self._make_instance_admin_api()
client = self._make_one(project=self.PROJECT, credentials=credentials)
client._instance_admin_api = api

Expand Down Expand Up @@ -754,15 +763,16 @@ def test_instance_factory_explicit(self):
self.assertIs(instance._client, client)

def test_list_instances(self):
from google.cloud.spanner_admin_instance_v1 import Instance as InstancePB
from google.cloud.spanner_admin_instance_v1 import (
InstanceAdminClient,
Instance as InstancePB,
)
from google.cloud.spanner_admin_instance_v1 import (
ListInstancesRequest,
ListInstancesResponse,
)

api = self._make_instance_admin_api()
credentials = build_scoped_credentials()
api = InstanceAdminClient(credentials=credentials)
client = self._make_one(project=self.PROJECT, credentials=credentials)
client._instance_admin_api = api

Expand Down Expand Up @@ -806,13 +816,12 @@ def test_list_instances(self):

def test_list_instances_w_options(self):
from google.cloud.spanner_admin_instance_v1 import (
InstanceAdminClient,
ListInstancesRequest,
ListInstancesResponse,
)

api = self._make_instance_admin_api()
credentials = build_scoped_credentials()
api = InstanceAdminClient(credentials=credentials)
client = self._make_one(project=self.PROJECT, credentials=credentials)
client._instance_admin_api = api

Expand Down
Loading
Loading