async def _status_code_and_message(stub: test_pb2_grpc.TestServiceStub):
details = 'test status message'
status = grpc.StatusCode.UNKNOWN # code = 2
# Test with a UnaryCall
request = messages_pb2.SimpleRequest(
response_type=messages_pb2.COMPRESSABLE,
response_size=1,
payload=messages_pb2.Payload(body=b'\x00'),
response_status=messages_pb2.EchoStatus(code=status.value[0],
message=details))
call = stub.UnaryCall(request)
await _validate_status_code_and_details(call, status, details)
# Test with a FullDuplexCall
call = stub.FullDuplexCall()
request = messages_pb2.StreamingOutputCallRequest(
response_type=messages_pb2.COMPRESSABLE,
response_parameters=(messages_pb2.ResponseParameters(size=1),),
payload=messages_pb2.Payload(body=b'\x00'),
response_status=messages_pb2.EchoStatus(code=status.value[0],
message=details))
await call.write(request) # sends the initial request.
await call.done_writing()
await _validate_status_code_and_details(call, status, details)
def _status_code_and_message(stub):
details = 'test status message'
code = 2
status = grpc.StatusCode.UNKNOWN # code = 2
# Test with a UnaryCall
request = messages_pb2.SimpleRequest(
response_type=messages_pb2.COMPRESSABLE,
response_size=1,
payload=messages_pb2.Payload(body=b'\x00'),
response_status=messages_pb2.EchoStatus(code=code, message=details))
response_future = stub.UnaryCall.future(request)
_validate_status_code_and_details(response_future, status, details)
# Test with a FullDuplexCall
with _Pipe() as pipe:
response_iterator = stub.FullDuplexCall(pipe)
request = messages_pb2.StreamingOutputCallRequest(
response_type=messages_pb2.COMPRESSABLE,
response_parameters=(messages_pb2.ResponseParameters(size=1), ),
payload=messages_pb2.Payload(body=b'\x00'),
response_status=messages_pb2.EchoStatus(code=code,
message=details))
pipe.add(request) # sends the initial request.
# Dropping out of with block closes the pipe
_validate_status_code_and_details(response_iterator, status, details)
def _ping_pong(stub):
request_response_sizes = (
31415,
9,
2653,
58979,
)
request_payload_sizes = (
27182,
8,
1828,
45904,
)
with _Pipe() as pipe:
response_iterator = stub.FullDuplexCall(pipe)
for response_size, payload_size in zip(request_response_sizes,
request_payload_sizes):
request = messages_pb2.StreamingOutputCallRequest(
response_type=messages_pb2.COMPRESSABLE,
response_parameters=(messages_pb2.ResponseParameters(
size=response_size), ),
payload=messages_pb2.Payload(body=b'\x00' * payload_size))
pipe.add(request)
response = next(response_iterator)
if response.payload.type != messages_pb2.COMPRESSABLE:
raise ValueError('response body of invalid type %s!' %
response.payload.type)
if len(response.payload.body) != response_size:
raise ValueError('response body of invalid size %d!' %
len(response.payload.body))
def StreamingOutputCall(self, request, context):
_maybe_echo_status_and_message(request, context)
for response_parameters in request.response_parameters:
yield messages_pb2.StreamingOutputCallResponse(
payload=messages_pb2.Payload(type=request.response_type,
body=b'\x00' *
response_parameters.size))
async def test_cancel_while_writing(self):
# Test cancelation before making any write or after doing at least 1
for num_writes_before_cancel in (0, 1):
with self.subTest(name="Num writes before cancel: {}".format(
num_writes_before_cancel)):
channel = aio.insecure_channel(
UNREACHABLE_TARGET,
interceptors=[_StreamUnaryInterceptorWithRequestIterator()])
stub = test_pb2_grpc.TestServiceStub(channel)
payload = messages_pb2.Payload(body=b'\0' *
_REQUEST_PAYLOAD_SIZE)
request = messages_pb2.StreamingInputCallRequest(
payload=payload)
call = stub.StreamingInputCall()
with self.assertRaises(asyncio.InvalidStateError):
for i in range(_NUM_STREAM_REQUESTS):
if i == num_writes_before_cancel:
self.assertTrue(call.cancel())
await call.write(request)
with self.assertRaises(asyncio.CancelledError):
await call
self.assertTrue(call.cancelled())
self.assertTrue(call.done())
self.assertEqual(await call.code(), grpc.StatusCode.CANCELLED)
await channel.close()
async def test_intercepts_request_iterator_rpc_error_using_write(self):
for interceptor_class in (_StreamUnaryInterceptorEmpty,
_StreamUnaryInterceptorWithRequestIterator):
with self.subTest(name=interceptor_class):
channel = aio.insecure_channel(
UNREACHABLE_TARGET, interceptors=[interceptor_class()])
stub = test_pb2_grpc.TestServiceStub(channel)
payload = messages_pb2.Payload(body=b'\0' *
_REQUEST_PAYLOAD_SIZE)
request = messages_pb2.StreamingInputCallRequest(
payload=payload)
call = stub.StreamingInputCall()
# When there is an error during the write, exception is raised.
with self.assertRaises(asyncio.InvalidStateError):
for _ in range(_NUM_STREAM_REQUESTS):
await call.write(request)
with self.assertRaises(aio.AioRpcError) as exception_context:
await call
self.assertEqual(grpc.StatusCode.UNAVAILABLE,
exception_context.exception.code())
self.assertTrue(call.done())
self.assertEqual(grpc.StatusCode.UNAVAILABLE, await call.code())
await channel.close()
async def test_add_done_callback_interceptor_task_finished(self):
for interceptor_class in (_StreamUnaryInterceptorEmpty,
_StreamUnaryInterceptorWithRequestIterator):
with self.subTest(name=interceptor_class):
interceptor = interceptor_class()
channel = aio.insecure_channel(self._server_target,
interceptors=[interceptor])
stub = test_pb2_grpc.TestServiceStub(channel)
payload = messages_pb2.Payload(body=b'\0' *
_REQUEST_PAYLOAD_SIZE)
request = messages_pb2.StreamingInputCallRequest(
payload=payload)
async def request_iterator():
for _ in range(_NUM_STREAM_REQUESTS):
yield request
call = stub.StreamingInputCall(request_iterator())
response = await call
validation = inject_callbacks(call)
await validation
await channel.close()
def __init__(self, server, config, hist):
# Create the stub
host, port = server.split(':')
port = int(port)
if config.HasField('security_params'):
creds = implementations.ssl_channel_credentials(
resources.test_root_certificates())
channel = test_utilities.not_really_secure_channel(
host, port, creds, config.security_params.server_host_override)
else:
channel = implementations.insecure_channel(host, port)
if config.payload_config.WhichOneof('payload') == 'simple_params':
self._generic = False
self._stub = services_pb2.beta_create_BenchmarkService_stub(
channel)
payload = messages_pb2.Payload(
body='\0' * config.payload_config.simple_params.req_size)
self._request = messages_pb2.SimpleRequest(
payload=payload,
response_size=config.payload_config.simple_params.resp_size)
else:
self._generic = True
self._stub = implementations.generic_stub(channel)
self._request = '\0' * config.payload_config.bytebuf_params.req_size
self._hist = hist
self._response_callbacks = []
async def test_interceptor_stream_stream(self):
record = []
server, stub = await _create_server_stub_pair(
_LoggingInterceptor('log_stream_stream', record))
# Prepares the request
payload = messages_pb2.Payload(body=b'\0' * _REQUEST_PAYLOAD_SIZE)
request = messages_pb2.StreamingInputCallRequest(payload=payload)
async def gen():
for _ in range(_NUM_STREAM_RESPONSES):
yield request
# Invokes the actual RPC
call = stub.StreamingInputCall(gen())
# Validates the responses
response = await call
self.assertIsInstance(response,
messages_pb2.StreamingInputCallResponse)
self.assertEqual(_NUM_STREAM_RESPONSES * _REQUEST_PAYLOAD_SIZE,
response.aggregated_payload_size)
self.assertEqual(await call.code(), grpc.StatusCode.OK)
self.assertSequenceEqual([
'log_stream_stream:intercept_service',
], record)
async def test_response_caching(self):
# Prepares a preset value to help testing
interceptor = _CacheInterceptor({
42:
messages_pb2.SimpleResponse(payload=messages_pb2.Payload(
body=b'\x42'))
})
# Constructs a server with the cache interceptor
server, stub = await _create_server_stub_pair(interceptor)
# Tests if the cache store is used
response = await stub.UnaryCall(
messages_pb2.SimpleRequest(response_size=42))
self.assertEqual(1, len(interceptor.cache_store[42].payload.body))
self.assertEqual(interceptor.cache_store[42], response)
# Tests response can be cached
response = await stub.UnaryCall(
messages_pb2.SimpleRequest(response_size=1337))
self.assertEqual(1337, len(interceptor.cache_store[1337].payload.body))
self.assertEqual(interceptor.cache_store[1337], response)
response = await stub.UnaryCall(
messages_pb2.SimpleRequest(response_size=1337))
self.assertEqual(interceptor.cache_store[1337], response)
def _cancel_after_first_response(stub):
request_response_sizes = (31415, 9, 2653, 58979,)
request_payload_sizes = (27182, 8, 1828, 45904,)
with _Pipe() as pipe:
response_iterator = stub.FullDuplexCall(pipe)
response_size = request_response_sizes[0]
payload_size = request_payload_sizes[0]
request = messages_pb2.StreamingOutputCallRequest(
response_type=messages_pb2.COMPRESSABLE,
response_parameters=(
messages_pb2.ResponseParameters(size=response_size),),
payload=messages_pb2.Payload(body=b'\x00' * payload_size))
pipe.add(request)
response = next(response_iterator)
# We test the contents of `response` in the Ping Pong test - don't check
# them here.
response_iterator.cancel()
try:
next(response_iterator)
except grpc.RpcError as rpc_error:
if rpc_error.code() is not grpc.StatusCode.CANCELLED:
raise
else:
raise ValueError('expected call to be cancelled')
def FullDuplexCall(self, request_iterator, context):
for request in request_iterator:
for response_parameters in request.response_parameters:
yield messages_pb2.StreamingOutputCallResponse(
payload=messages_pb2.Payload(type=request.payload.type,
body=b'\x00' *
response_parameters.size))
async def test_stream_unary_using_write(self):
channel = aio.insecure_channel(self._server_target)
stub = test_pb2_grpc.TestServiceStub(channel)
# Invokes the actual RPC
call = stub.StreamingInputCall()
# Prepares the request
payload = messages_pb2.Payload(body=b'\0' * _REQUEST_PAYLOAD_SIZE)
request = messages_pb2.StreamingInputCallRequest(payload=payload)
# Sends out requests
for _ in range(_NUM_STREAM_RESPONSES):
await call.write(request)
await call.done_writing()
# Validates the responses
response = await call
self.assertIsInstance(response,
messages_pb2.StreamingInputCallResponse)
self.assertEqual(_NUM_STREAM_RESPONSES * _REQUEST_PAYLOAD_SIZE,
response.aggregated_payload_size)
self.assertEqual(await call.code(), grpc.StatusCode.OK)
await channel.close()
def UnaryCall(self, request, context):
if request.HasField('response_status'):
context.set_code(request.response_status.code)
context.set_details(request.response_status.message)
return messages_pb2.SimpleResponse(
payload=messages_pb2.Payload(type=messages_pb2.COMPRESSABLE,
body=b'\x00' * request.response_size))
def _status_code_and_message(stub):
message = 'test status message'
code = 2
status = grpc.StatusCode.UNKNOWN # code = 2
request = messages_pb2.SimpleRequest(
response_type=messages_pb2.COMPRESSABLE,
response_size=1,
payload=messages_pb2.Payload(body=b'\x00'),
response_status=messages_pb2.EchoStatus(code=code, message=message))
response_future = stub.UnaryCall.future(request)
if response_future.code() != status:
raise ValueError('expected code %s, got %s' %
(status, response_future.code()))
elif response_future.details() != message:
raise ValueError('expected message %s, got %s' %
(message, response_future.details()))
request = messages_pb2.StreamingOutputCallRequest(
response_type=messages_pb2.COMPRESSABLE,
response_parameters=(messages_pb2.ResponseParameters(size=1), ),
response_status=messages_pb2.EchoStatus(code=code, message=message))
response_iterator = stub.StreamingOutputCall(request)
if response_future.code() != status:
raise ValueError('expected code %s, got %s' %
(status, response_iterator.code()))
elif response_future.details() != message:
raise ValueError('expected message %s, got %s' %
(message, response_iterator.details()))
def UnaryCall(self, request, context):
_maybe_echo_metadata(context)
_maybe_echo_status_and_message(request, context)
return messages_pb2.SimpleResponse(
payload=messages_pb2.Payload(
type=messages_pb2.COMPRESSABLE,
body=b'\x00' * request.response_size))
async def test_cancel_by_the_interceptor(self):
class Interceptor(aio.StreamUnaryClientInterceptor):
async def intercept_stream_unary(self, continuation,
client_call_details,
request_iterator):
call = await continuation(client_call_details, request_iterator)
call.cancel()
return call
channel = aio.insecure_channel(UNREACHABLE_TARGET,
interceptors=[Interceptor()])
stub = test_pb2_grpc.TestServiceStub(channel)
payload = messages_pb2.Payload(body=b'\0' * _REQUEST_PAYLOAD_SIZE)
request = messages_pb2.StreamingInputCallRequest(payload=payload)
call = stub.StreamingInputCall()
with self.assertRaises(asyncio.InvalidStateError):
for i in range(_NUM_STREAM_REQUESTS):
await call.write(request)
with self.assertRaises(asyncio.CancelledError):
await call
self.assertTrue(call.cancelled())
self.assertTrue(call.done())
self.assertEqual(await call.code(), grpc.StatusCode.CANCELLED)
await channel.close()
def _ping_pong(stub):
request_response_sizes = (
31415,
9,
2653,
58979,
)
request_payload_sizes = (
27182,
8,
1828,
45904,
)
with _Pipe() as pipe:
response_iterator = stub.FullDuplexCall(pipe)
for response_size, payload_size in zip(request_response_sizes,
request_payload_sizes):
request = messages_pb2.StreamingOutputCallRequest(
response_type=messages_pb2.COMPRESSABLE,
response_parameters=(messages_pb2.ResponseParameters(
size=response_size), ),
payload=messages_pb2.Payload(body=b'\x00' * payload_size))
pipe.add(request)
response = next(response_iterator)
_validate_payload_type_and_length(response,
messages_pb2.COMPRESSABLE,
response_size)
async def test_exception_raised_by_interceptor(self):
class InterceptorException(Exception):
pass
class Interceptor(aio.StreamUnaryClientInterceptor):
async def intercept_stream_unary(self, continuation,
client_call_details,
request_iterator):
raise InterceptorException
channel = aio.insecure_channel(UNREACHABLE_TARGET,
interceptors=[Interceptor()])
stub = test_pb2_grpc.TestServiceStub(channel)
payload = messages_pb2.Payload(body=b'\0' * _REQUEST_PAYLOAD_SIZE)
request = messages_pb2.StreamingInputCallRequest(payload=payload)
call = stub.StreamingInputCall()
with self.assertRaises(InterceptorException):
for i in range(_NUM_STREAM_REQUESTS):
await call.write(request)
with self.assertRaises(InterceptorException):
await call
await channel.close()
def _cancel_after_first_response(stub):
request_response_sizes = (31415, 9, 2653, 58979)
request_payload_sizes = (27182, 8, 1828, 45904)
with stub, _Pipe() as pipe:
response_iterator = stub.FullDuplexCall(pipe, _TIMEOUT)
response_size = request_response_sizes[0]
payload_size = request_payload_sizes[0]
request = messages_pb2.StreamingOutputCallRequest(
response_type=messages_pb2.COMPRESSABLE,
response_parameters=(messages_pb2.ResponseParameters(
size=response_size), ),
payload=messages_pb2.Payload(body=b'\x00' * payload_size))
pipe.add(request)
response = next(response_iterator)
# We test the contents of `response` in the Ping Pong test - don't check
# them here.
response_iterator.cancel()
try:
next(response_iterator)
except Exception:
pass
else:
raise ValueError('expected call to be cancelled')
def __init__(self, server, config, hist):
# Create the stub
if config.HasField('security_params'):
creds = grpc.ssl_channel_credentials(
resources.test_root_certificates())
channel = test_common.test_secure_channel(
server, creds, config.security_params.server_host_override)
else:
channel = grpc.insecure_channel(server)
# waits for the channel to be ready before we start sending messages
grpc.channel_ready_future(channel).result()
if config.payload_config.WhichOneof('payload') == 'simple_params':
self._generic = False
self._stub = benchmark_service_pb2_grpc.BenchmarkServiceStub(
channel)
payload = messages_pb2.Payload(
body=bytes(b'\0' *
config.payload_config.simple_params.req_size))
self._request = messages_pb2.SimpleRequest(
payload=payload,
response_size=config.payload_config.simple_params.resp_size)
else:
self._generic = True
self._stub = GenericStub(channel)
self._request = bytes(b'\0' *
config.payload_config.bytebuf_params.req_size)
self._hist = hist
self._response_callbacks = []
async def _cancel_after_first_response(stub: test_pb2_grpc.TestServiceStub):
request_response_sizes = (
31415,
9,
2653,
58979,
)
request_payload_sizes = (
27182,
8,
1828,
45904,
)
call = stub.FullDuplexCall()
response_size = request_response_sizes[0]
payload_size = request_payload_sizes[0]
request = messages_pb2.StreamingOutputCallRequest(
response_type=messages_pb2.COMPRESSABLE,
response_parameters=(messages_pb2.ResponseParameters(
size=response_size),),
payload=messages_pb2.Payload(body=b'\x00' * payload_size))
await call.write(request)
await call.read()
call.cancel()
try:
await call.read()
except asyncio.CancelledError:
assert await call.code() is grpc.StatusCode.CANCELLED
else:
raise ValueError('expected call to be cancelled')
async def _custom_metadata(stub: test_pb2_grpc.TestServiceStub):
initial_metadata_value = "test_initial_metadata_value"
trailing_metadata_value = b"\x0a\x0b\x0a\x0b\x0a\x0b"
metadata = ((_INITIAL_METADATA_KEY, initial_metadata_value),
(_TRAILING_METADATA_KEY, trailing_metadata_value))
async def _validate_metadata(call):
initial_metadata = dict(await call.initial_metadata())
if initial_metadata[_INITIAL_METADATA_KEY] != initial_metadata_value:
raise ValueError('expected initial metadata %s, got %s' %
(initial_metadata_value,
initial_metadata[_INITIAL_METADATA_KEY]))
trailing_metadata = dict(await call.trailing_metadata())
if trailing_metadata[_TRAILING_METADATA_KEY] != trailing_metadata_value:
raise ValueError('expected trailing metadata %s, got %s' %
(trailing_metadata_value,
trailing_metadata[_TRAILING_METADATA_KEY]))
# Testing with UnaryCall
request = messages_pb2.SimpleRequest(
response_type=messages_pb2.COMPRESSABLE,
response_size=1,
payload=messages_pb2.Payload(body=b'\x00'))
call = stub.UnaryCall(request, metadata=metadata)
await _validate_metadata(call)
# Testing with FullDuplexCall
call = stub.FullDuplexCall(metadata=metadata)
request = messages_pb2.StreamingOutputCallRequest(
response_type=messages_pb2.COMPRESSABLE,
response_parameters=(messages_pb2.ResponseParameters(size=1),))
await call.write(request)
await call.read()
await call.done_writing()
await _validate_metadata(call)
async def _ping_pong(stub: test_pb2_grpc.TestServiceStub) -> None:
request_response_sizes = (
31415,
9,
2653,
58979,
)
request_payload_sizes = (
27182,
8,
1828,
45904,
)
call = stub.FullDuplexCall()
for response_size, payload_size in zip(request_response_sizes,
request_payload_sizes):
request = messages_pb2.StreamingOutputCallRequest(
response_type=messages_pb2.COMPRESSABLE,
response_parameters=(messages_pb2.ResponseParameters(
size=response_size),),
payload=messages_pb2.Payload(body=b'\x00' * payload_size))
await call.write(request)
response = await call.read()
_validate_payload_type_and_length(response, messages_pb2.COMPRESSABLE,
response_size)
await call.done_writing()
await _validate_status_code_and_details(call, grpc.StatusCode.OK, '')
def __init__(self, server, config, hist):
# Create the stub
if config.HasField('security_params'):
creds = grpc.ssl_channel_credentials(resources.test_root_certificates())
channel = test_common.test_secure_channel(
server, creds, config.security_params.server_host_override)
else:
channel = grpc.insecure_channel(server)
connected_event = threading.Event()
def wait_for_ready(connectivity):
if connectivity == grpc.ChannelConnectivity.READY:
connected_event.set()
channel.subscribe(wait_for_ready, try_to_connect=True)
connected_event.wait()
if config.payload_config.WhichOneof('payload') == 'simple_params':
self._generic = False
self._stub = services_pb2.BenchmarkServiceStub(channel)
payload = messages_pb2.Payload(
body='\0' * config.payload_config.simple_params.req_size)
self._request = messages_pb2.SimpleRequest(
payload=payload,
response_size=config.payload_config.simple_params.resp_size)
else:
self._generic = True
self._stub = GenericStub(channel)
self._request = '\0' * config.payload_config.bytebuf_params.req_size
self._hist = hist
self._response_callbacks = []
def _blocking_unary(stub):
size = 314159
request = messages_pb2.SimpleRequest(
response_type=messages_pb2.COMPRESSABLE,
response_size=size,
payload=messages_pb2.Payload(body=b'\x00' * 271828))
response = stub.UnaryCall(request, timeout=_RPC_TIMEOUT_S)
_validate_payload_type_and_length(response, messages_pb2.COMPRESSABLE, size)
def StreamingOutputCall(self, request, context):
if request.HasField('response_status'):
context.set_code(request.response_status.code)
context.set_details(request.response_status.message)
for response_parameters in request.response_parameters:
yield messages_pb2.StreamingOutputCallResponse(
payload=messages_pb2.Payload(type=request.response_type,
body=b'\x00' *
response_parameters.size))
def StreamingOutputCall(self, request, context):
_maybe_echo_status_and_message(request, context)
for response_parameters in request.response_parameters:
if response_parameters.interval_us != 0:
time.sleep(response_parameters.interval_us / _US_IN_A_SECOND)
yield messages_pb2.StreamingOutputCallResponse(
payload=messages_pb2.Payload(type=request.response_type,
body=b'\x00' *
response_parameters.size))
def FullDuplexCall(self, request_iterator, context):
_maybe_echo_metadata(context)
for request in request_iterator:
_maybe_echo_status_and_message(request, context)
for response_parameters in request.response_parameters:
yield messages_pb2.StreamingOutputCallResponse(
payload=messages_pb2.Payload(type=request.payload.type,
body=b'\x00' *
response_parameters.size))
def _cancel_after_begin(stub):
sizes = (27182, 8, 1828, 45904,)
payloads = (messages_pb2.Payload(body=b'\x00' * size) for size in sizes)
requests = (messages_pb2.StreamingInputCallRequest(payload=payload)
for payload in payloads)
response_future = stub.StreamingInputCall.future(requests)
response_future.cancel()
if not response_future.cancelled():
raise ValueError('expected call to be cancelled')
