Skip to content
This repository was archived by the owner on Nov 12, 2025. It is now read-only.

chore: test that ServiceUnavailable is retried #970

Merged
tswast merged 2 commits into from
Aug 14, 2025
Merged

chore: test that ServiceUnavailable is retried #970

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
87b8358
chore: test that ServiceUnavailable is retried
tswast Aug 12, 2025
3187b42
remove debug file
tswast Aug 13, 2025
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
13 changes: 13 additions & 0 deletions main.py
View file
Open in desktop
Comment thread
tswast marked this conversation as resolved.
Outdated
Original file line number Diff line number Diff line change
@@ -0,0 +1,13 @@

import google.cloud.bigquery_storage_v1
import google.cloud.bigquery_storage_v1.types

readclient = google.cloud.bigquery_storage_v1.BigQueryReadClient()
request = google.cloud.bigquery_storage_v1.types.CreateReadSessionRequest(
parent="projects/swast-scratch",
read_session=google.cloud.bigquery_storage_v1.types.ReadSession(
table="projects/swast-scratch/datasets/my_dataset/tables/all_types",
data_format=google.cloud.bigquery_storage_v1.types.DataFormat.ARROW,
))

session = readclient.create_read_session(request)
56 changes: 45 additions & 11 deletions tests/unit/test_client_v1.py → tests/unit/test_read_client_v1.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -15,11 +15,12 @@
import importlib
from unittest import mock

import google.api_core.exceptions
from google.api_core.gapic_v1 import client_info
from google.auth import credentials
import pytest

from google.cloud.bigquery_storage import types
from google.cloud.bigquery_storage_v1 import types

PROJECT = "my-project"
SERVICE_ACCOUNT_PROJECT = "project-from-credentials"
Expand All @@ -29,20 +30,15 @@
def mock_transport(monkeypatch):
from google.cloud.bigquery_storage_v1.services.big_query_read import transports

fake_create_session_rpc = mock.Mock(name="create_read_session_rpc")
fake_read_rows_rpc = mock.Mock(name="read_rows_rpc")

transport = mock.create_autospec(
transports.grpc.BigQueryReadGrpcTransport, instance=True
)

transport.create_read_session = mock.Mock(name="fake_create_read_session")
transport.read_rows = mock.Mock(name="fake_read_rows")

transport._wrapped_methods = {
transport.create_read_session: fake_create_session_rpc,
transport.read_rows: fake_read_rows_rpc,
}
transports.grpc.BigQueryReadGrpcTransport._prep_wrapped_messages(
transport, client_info.ClientInfo()
)

# _credentials property for TPC support
transport._credentials = ""
Expand Down Expand Up @@ -85,8 +81,11 @@ def __init__(self, *args, **kwargs):


def test_create_read_session(mock_transport, client_under_test):
assert client_under_test._transport is mock_transport # sanity check
# validate test assumptions
assert client_under_test._transport is mock_transport

rpc_callable = mock.Mock()
mock_transport._wrapped_methods[mock_transport.create_read_session] = rpc_callable
table = "projects/{}/datasets/{}/tables/{}".format(
"data-project-id", "dataset_id", "table_id"
)
Expand All @@ -101,12 +100,47 @@ def test_create_read_session(mock_transport, client_under_test):
expected_session_arg = types.CreateReadSessionRequest(
parent="projects/other-project", read_session=read_session
)
rpc_callable = mock_transport._wrapped_methods[mock_transport.create_read_session]
rpc_callable.assert_called_once_with(
expected_session_arg, metadata=mock.ANY, retry=mock.ANY, timeout=mock.ANY
)


def test_create_read_session_retries_serviceunavailable(
mock_transport, client_under_test
):
"""Regression test for https://github.com/googleapis/python-bigquery-storage/issues/969."""
# validate test assumptions
assert client_under_test._transport is mock_transport

mock_transport.create_read_session.side_effect = [

@chalmerlowe chalmerlowe Aug 14, 2025 •
edited
Loading

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@tswast

QUESTION:

When we assign a list to the attribute class.method.side_effect (i.e. mock_transport.create_read_session.side_effect), every time the method gets called, the next item in the list will be produced as the result of the method call.

We have three items in the list.

One line 135 we confirm that the create_read_session method has been called thrice.
But I do not see where we make those three calls.
I see one call on line 127.

What am I missing?

Besides this question, everything else LGTM and will approve on that basis.

@tswast tswast Aug 14, 2025

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

But I do not see where we make those three calls.

There are several layers in the client library. To write this test, I had to figure out how to mock out the level that gets wrapped by the retry logic, which is about as low of a layer we can get in Python before it switches over to gRPC C++ stuff. The retry method is implemented as a wrapper on top of this method. So the 3 calls are my way of verifying that the first two get retried.

google.api_core.exceptions.ServiceUnavailable("connection reset"),
google.api_core.exceptions.ServiceUnavailable("connection reset"),
types.ReadSession(),
]
table = "projects/{}/datasets/{}/tables/{}".format(
"data-project-id", "dataset_id", "table_id"
)
read_session = types.ReadSession()
read_session.table = table

# with pytest.raises(google.api_core.exceptions.ServiceUnavailable):
client_under_test.create_read_session(
parent="projects/other-project", read_session=read_session
)

expected_session_arg = types.CreateReadSessionRequest(
parent="projects/other-project", read_session=read_session
)
expected_call = mock.call(expected_session_arg, metadata=mock.ANY, timeout=mock.ANY)
mock_transport.create_read_session.assert_has_calls(
[
expected_call,
expected_call,
expected_call,
]
)


def test_read_rows(mock_transport, client_under_test):
stream_name = "teststream"
offset = 0
Expand Down