def test_create_read_session(self):
# Setup Expected Response
name = "name3373707"
table = "table110115790"
expected_response = {"name": name, "table": table}
expected_response = stream_pb2.ReadSession(**expected_response)
# Mock the API response
channel = ChannelStub(responses=[expected_response])
patch = mock.patch("google.api_core.grpc_helpers.create_channel")
with patch as create_channel:
create_channel.return_value = channel
client = big_query_read_client.BigQueryReadClient()
# Setup Request
parent = client.project_path("[PROJECT]")
read_session = {}
response = client.create_read_session(parent, read_session)
assert expected_response == response
assert len(channel.requests) == 1
expected_request = storage_pb2.CreateReadSessionRequest(
parent=parent, read_session=read_session
)
actual_request = channel.requests[0][1]
assert expected_request == actual_request
def test_read_rows(self):
# Setup Expected Response
row_count = 1340416618
expected_response = {"row_count": row_count}
expected_response = storage_pb2.ReadRowsResponse(**expected_response)
# Mock the API response
channel = ChannelStub(responses=[iter([expected_response])])
patch = mock.patch("google.api_core.grpc_helpers.create_channel")
with patch as create_channel:
create_channel.return_value = channel
client = big_query_read_client.BigQueryReadClient()
# Setup Request
read_stream = client.read_stream_path(
"[PROJECT]", "[LOCATION]", "[SESSION]", "[STREAM]"
)
response = client.read_rows(read_stream)
resources = list(response)
assert len(resources) == 1
assert expected_response == resources[0]
assert len(channel.requests) == 1
expected_request = storage_pb2.ReadRowsRequest(read_stream=read_stream)
actual_request = channel.requests[0][1]
assert expected_request == actual_request
def test_create_read_session_exception(self):
# Mock the API response
channel = ChannelStub(responses=[CustomException()])
patch = mock.patch("google.api_core.grpc_helpers.create_channel")
with patch as create_channel:
create_channel.return_value = channel
client = big_query_read_client.BigQueryReadClient()
with pytest.raises(CustomException):
client.create_read_session()
def test_split_read_stream_exception(self):
# Mock the API response
channel = ChannelStub(responses=[CustomException()])
patch = mock.patch("google.api_core.grpc_helpers.create_channel")
with patch as create_channel:
create_channel.return_value = channel
client = big_query_read_client.BigQueryReadClient()
# Setup request
name = client.read_stream_path(
"[PROJECT]", "[LOCATION]", "[SESSION]", "[STREAM]"
)
with pytest.raises(CustomException):
client.split_read_stream(name)
def test_split_read_stream(self):
# Setup Expected Response
expected_response = {}
expected_response = storage_pb2.SplitReadStreamResponse(
**expected_response)
# Mock the API response
channel = ChannelStub(responses=[expected_response])
patch = mock.patch("google.api_core.grpc_helpers.create_channel")
with patch as create_channel:
create_channel.return_value = channel
client = big_query_read_client.BigQueryReadClient()
response = client.split_read_stream()
assert expected_response == response
assert len(channel.requests) == 1
expected_request = storage_pb2.SplitReadStreamRequest()
actual_request = channel.requests[0][1]
assert expected_request == actual_request
