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_storage_client.BigQueryStorageClient()
# Setup request
original_stream = {}
with pytest.raises(CustomException):
client.split_read_stream(original_stream)
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_storage_client.BigQueryStorageClient()
# Setup request
table_reference = {}
parent = "parent-995424086"
with pytest.raises(CustomException):
client.create_read_session(table_reference, parent)
def test_batch_create_read_session_streams_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_storage_client.BigQueryStorageClient()
# Setup request
session = {}
requested_streams = 1017221410
with pytest.raises(CustomException):
client.batch_create_read_session_streams(session, requested_streams)
def test_finalize_stream(self):
channel = ChannelStub()
patch = mock.patch("google.api_core.grpc_helpers.create_channel")
with patch as create_channel:
create_channel.return_value = channel
client = big_query_storage_client.BigQueryStorageClient()
# Setup Request
stream = {}
client.finalize_stream(stream)
assert len(channel.requests) == 1
expected_request = storage_pb2.FinalizeStreamRequest(stream=stream)
actual_request = channel.requests[0][1]
assert expected_request == actual_request
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_storage_client.BigQueryStorageClient()
# Setup Request
original_stream = {}
response = client.split_read_stream(original_stream)
assert expected_response == response
assert len(channel.requests) == 1
expected_request = storage_pb2.SplitReadStreamRequest(
original_stream=original_stream
)
actual_request = channel.requests[0][1]
assert expected_request == actual_request
def test_read_rows(self):
# Setup Expected Response
expected_response = {}
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_storage_client.BigQueryStorageClient()
# Setup Request
read_position = {}
response = client.read_rows(read_position)
resources = list(response)
assert len(resources) == 1
assert expected_response == resources[0]
assert len(channel.requests) == 1
expected_request = storage_pb2.ReadRowsRequest(read_position=read_position)
actual_request = channel.requests[0][1]
assert expected_request == actual_request
