async def setUp(self):
address, self._async_server = await start_test_server()
# Create async stub
self._async_channel = aio.insecure_channel(address,
options=_unique_options())
self._async_stub = test_pb2_grpc.TestServiceStub(self._async_channel)
# Create sync stub
self._sync_channel = grpc.insecure_channel(address,
options=_unique_options())
self._sync_stub = test_pb2_grpc.TestServiceStub(self._sync_channel)
async def test_channel_isolation(self):
async with aio.insecure_channel(self._server_target) as channel1:
async with aio.insecure_channel(self._server_target) as channel2:
stub1 = test_pb2_grpc.TestServiceStub(channel1)
stub2 = test_pb2_grpc.TestServiceStub(channel2)
call1 = stub1.UnaryCall(messages_pb2.SimpleRequest())
call2 = stub2.UnaryCall(messages_pb2.SimpleRequest())
self.assertFalse(call1.cancelled())
self.assertTrue(call2.cancelled())
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()
async def test_initial_metadata_modification(self):
class Interceptor(aio.UnaryUnaryClientInterceptor):
async def intercept_unary_unary(self, continuation,
client_call_details, request):
new_details = aio.ClientCallDetails(
method=client_call_details.method,
timeout=client_call_details.timeout,
metadata=client_call_details.metadata +
_INITIAL_METADATA_TO_INJECT,
credentials=client_call_details.credentials,
wait_for_ready=client_call_details.wait_for_ready,
)
return await continuation(new_details, request)
async with aio.insecure_channel(self._server_target,
interceptors=[Interceptor()
]) as channel:
stub = test_pb2_grpc.TestServiceStub(channel)
call = stub.UnaryCall(messages_pb2.SimpleRequest())
# Expected to see the echoed initial metadata
self.assertTrue(
_common.seen_metadatum(_INITIAL_METADATA_TO_INJECT[0], await
call.initial_metadata()))
# Expected to see the echoed trailing metadata
self.assertTrue(
_common.seen_metadatum(_INITIAL_METADATA_TO_INJECT[1], await
call.trailing_metadata()))
self.assertEqual(await call.code(), grpc.StatusCode.OK)
def _run_single_channel(config: _ChannelConfiguration) -> None:
global _global_rpc_id # pylint: disable=global-statement
with config.condition:
server = config.server
with grpc.insecure_channel(server) as channel:
stub = test_pb2_grpc.TestServiceStub(channel)
futures: Dict[int, Tuple[grpc.Future, str]] = {}
while not _stop_event.is_set():
with config.condition:
if config.qps == 0:
config.condition.wait(
timeout=_CONFIG_CHANGE_TIMEOUT.total_seconds())
continue
else:
duration_per_query = 1.0 / float(config.qps)
request_id = None
with _global_lock:
request_id = _global_rpc_id
_global_rpc_id += 1
_global_rpcs_started[config.method] += 1
start = time.time()
end = start + duration_per_query
with config.condition:
_start_rpc(config.method, config.metadata, request_id, stub,
float(config.rpc_timeout_sec), futures)
with config.condition:
_remove_completed_rpcs(futures, config.print_response)
logger.debug(f"Currently {len(futures)} in-flight RPCs")
now = time.time()
while now < end:
time.sleep(end - now)
now = time.time()
_cancel_all_rpcs(futures)
async def test_unary_stream(self):
channel = aio.insecure_channel(self._server_target)
stub = test_pb2_grpc.TestServiceStub(channel)
# Prepares the request
request = messages_pb2.StreamingOutputCallRequest()
for _ in range(_NUM_STREAM_RESPONSES):
request.response_parameters.append(
messages_pb2.ResponseParameters(size=_RESPONSE_PAYLOAD_SIZE))
# Invokes the actual RPC
call = stub.StreamingOutputCall(request)
# Validates the responses
response_cnt = 0
async for response in call:
response_cnt += 1
self.assertIs(type(response),
messages_pb2.StreamingOutputCallResponse)
self.assertEqual(_RESPONSE_PAYLOAD_SIZE,
len(response.payload.body))
self.assertEqual(_NUM_STREAM_RESPONSES, response_cnt)
self.assertEqual(await call.code(), grpc.StatusCode.OK)
await channel.close()
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()
async def test_late_cancel_unary_stream(self):
"""Test cancellation after received all messages."""
async with aio.insecure_channel(self._server_target) as channel:
stub = test_pb2_grpc.TestServiceStub(channel)
# Prepares the request
request = messages_pb2.StreamingOutputCallRequest()
for _ in range(_NUM_STREAM_RESPONSES):
request.response_parameters.append(
messages_pb2.ResponseParameters(
size=_RESPONSE_PAYLOAD_SIZE,))
# Invokes the actual RPC
call = stub.StreamingOutputCall(request)
for _ in range(_NUM_STREAM_RESPONSES):
response = await call.read()
self.assertIs(type(response),
messages_pb2.StreamingOutputCallResponse)
self.assertEqual(_RESPONSE_PAYLOAD_SIZE,
len(response.payload.body))
# After all messages received, it is possible that the final state
# is received or on its way. It's basically a data race, so our
# expectation here is do not crash :)
call.cancel()
self.assertIn(await call.code(),
[grpc.StatusCode.OK, grpc.StatusCode.CANCELLED])
async def test_cancel_after_rpc(self):
interceptor_reached = asyncio.Event()
wait_for_ever = self.loop.create_future()
class Interceptor(aio.UnaryStreamClientInterceptor):
async def intercept_unary_stream(self, continuation,
client_call_details, request):
call = await continuation(client_call_details, request)
interceptor_reached.set()
await wait_for_ever
channel = aio.insecure_channel(UNREACHABLE_TARGET,
interceptors=[Interceptor()])
request = messages_pb2.StreamingOutputCallRequest()
stub = test_pb2_grpc.TestServiceStub(channel)
call = stub.StreamingOutputCall(request)
self.assertFalse(call.cancelled())
self.assertFalse(call.done())
await interceptor_reached.wait()
self.assertTrue(call.cancel())
with self.assertRaises(asyncio.CancelledError):
async for response in call:
pass
self.assertTrue(call.cancelled())
self.assertTrue(call.done())
self.assertEqual(await call.code(), grpc.StatusCode.CANCELLED)
self.assertEqual(await call.initial_metadata(), None)
self.assertEqual(await call.trailing_metadata(), None)
await channel.close()
async def test_multiple_cancel_unary_stream(self):
async with aio.insecure_channel(self._server_target) as channel:
stub = test_pb2_grpc.TestServiceStub(channel)
# Prepares the request
request = messages_pb2.StreamingOutputCallRequest()
for _ in range(_NUM_STREAM_RESPONSES):
request.response_parameters.append(
messages_pb2.ResponseParameters(
size=_RESPONSE_PAYLOAD_SIZE,
interval_us=_RESPONSE_INTERVAL_US,
))
# Invokes the actual RPC
call = stub.StreamingOutputCall(request)
self.assertFalse(call.cancelled())
response = await call.read()
self.assertIs(type(response),
messages_pb2.StreamingOutputCallResponse)
self.assertEqual(_RESPONSE_PAYLOAD_SIZE, len(response.payload.body))
self.assertTrue(call.cancel())
self.assertFalse(call.cancel())
self.assertFalse(call.cancel())
self.assertFalse(call.cancel())
with self.assertRaises(asyncio.CancelledError):
await call.read()
async def test_response_iterator_using_read(self):
interceptor = _UnaryStreamInterceptorWithResponseIterator()
channel = aio.insecure_channel(self._server_target,
interceptors=[interceptor])
stub = test_pb2_grpc.TestServiceStub(channel)
request = messages_pb2.StreamingOutputCallRequest()
request.response_parameters.extend(
[messages_pb2.ResponseParameters(size=_RESPONSE_PAYLOAD_SIZE)] *
_NUM_STREAM_RESPONSES)
call = stub.StreamingOutputCall(request)
response_cnt = 0
for response in range(_NUM_STREAM_RESPONSES):
response = await call.read()
response_cnt += 1
self.assertIs(type(response),
messages_pb2.StreamingOutputCallResponse)
self.assertEqual(_RESPONSE_PAYLOAD_SIZE,
len(response.payload.body))
self.assertEqual(response_cnt, _NUM_STREAM_RESPONSES)
self.assertEqual(interceptor.response_iterator.response_cnt,
_NUM_STREAM_RESPONSES)
self.assertEqual(await call.code(), grpc.StatusCode.OK)
await channel.close()
async def test_multiple_interceptors_response_iterator(self):
for interceptor_class in (_UnaryStreamInterceptorEmpty,
_UnaryStreamInterceptorWithResponseIterator):
with self.subTest(name=interceptor_class):
interceptors = [interceptor_class(), interceptor_class()]
channel = aio.insecure_channel(self._server_target,
interceptors=interceptors)
stub = test_pb2_grpc.TestServiceStub(channel)
request = messages_pb2.StreamingOutputCallRequest()
request.response_parameters.extend([
messages_pb2.ResponseParameters(
size=_RESPONSE_PAYLOAD_SIZE)
] * _NUM_STREAM_RESPONSES)
call = stub.StreamingOutputCall(request)
response_cnt = 0
async for response in call:
response_cnt += 1
self.assertIs(type(response),
messages_pb2.StreamingOutputCallResponse)
self.assertEqual(_RESPONSE_PAYLOAD_SIZE,
len(response.payload.body))
self.assertEqual(response_cnt, _NUM_STREAM_RESPONSES)
self.assertEqual(await call.code(), grpc.StatusCode.OK)
await channel.close()
async def test_cancel_unary_stream_in_task_using_async_for(self):
async with aio.insecure_channel(self._server_target) as channel:
stub = test_pb2_grpc.TestServiceStub(channel)
coro_started = asyncio.Event()
# Configs the server method to block forever
request = messages_pb2.StreamingOutputCallRequest()
request.response_parameters.append(
messages_pb2.ResponseParameters(
size=_RESPONSE_PAYLOAD_SIZE,
interval_us=_INFINITE_INTERVAL_US,
))
# Invokes the actual RPC
call = stub.StreamingOutputCall(request)
async def another_coro():
coro_started.set()
async for _ in call:
pass
task = self.loop.create_task(another_coro())
await coro_started.wait()
self.assertFalse(task.done())
task.cancel()
self.assertEqual(grpc.StatusCode.CANCELLED, await call.code())
with self.assertRaises(asyncio.CancelledError):
await task
async def test_add_done_callback_interceptor_task_finished(self):
for interceptor_class in (_UnaryStreamInterceptorEmpty,
_UnaryStreamInterceptorWithResponseIterator):
with self.subTest(name=interceptor_class):
interceptor = interceptor_class()
request = messages_pb2.StreamingOutputCallRequest()
request.response_parameters.extend([
messages_pb2.ResponseParameters(
size=_RESPONSE_PAYLOAD_SIZE)
] * _NUM_STREAM_RESPONSES)
channel = aio.insecure_channel(self._server_target,
interceptors=[interceptor])
stub = test_pb2_grpc.TestServiceStub(channel)
call = stub.StreamingOutputCall(request)
# This ensures that the callbacks will be registered
# with the intercepted call rather than saving in the
# pending state list.
await call.wait_for_connection()
validation = inject_callbacks(call)
async for response in call:
pass
await validation
await channel.close()
async def test_max_message_length_applied(self):
address, server = await start_test_server()
async with aio.insecure_channel(
address,
options=((_GRPC_ARG_MAX_RECEIVE_MESSAGE_LENGTH,
_MAX_MESSAGE_LENGTH),)) as channel:
stub = test_pb2_grpc.TestServiceStub(channel)
request = messages_pb2.StreamingOutputCallRequest()
# First request will pass
request.response_parameters.append(
messages_pb2.ResponseParameters(size=_MAX_MESSAGE_LENGTH // 2,))
# Second request should fail
request.response_parameters.append(
messages_pb2.ResponseParameters(size=_MAX_MESSAGE_LENGTH * 2,))
call = stub.StreamingOutputCall(request)
response = await call.read()
self.assertEqual(_MAX_MESSAGE_LENGTH // 2,
len(response.payload.body))
with self.assertRaises(aio.AioRpcError) as exception_context:
await call.read()
rpc_error = exception_context.exception
self.assertEqual(grpc.StatusCode.RESOURCE_EXHAUSTED,
rpc_error.code())
self.assertIn(str(_MAX_MESSAGE_LENGTH), rpc_error.details())
self.assertEqual(grpc.StatusCode.RESOURCE_EXHAUSTED, await
call.code())
await server.stop(None)
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()
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 test_graceful_close(self):
stub = test_pb2_grpc.TestServiceStub(self._channel)
_, response = stub.UnaryCall.with_call(messages_pb2.SimpleRequest())
self._channel.close()
self.assertEqual(grpc.StatusCode.OK, response.code())
async def test_too_many_reads_unary_stream(self):
"""Test cancellation after received all messages."""
async with aio.insecure_channel(self._server_target) as channel:
stub = test_pb2_grpc.TestServiceStub(channel)
# Prepares the request
request = messages_pb2.StreamingOutputCallRequest()
for _ in range(_NUM_STREAM_RESPONSES):
request.response_parameters.append(
messages_pb2.ResponseParameters(
size=_RESPONSE_PAYLOAD_SIZE,))
# Invokes the actual RPC
call = stub.StreamingOutputCall(request)
for _ in range(_NUM_STREAM_RESPONSES):
response = await call.read()
self.assertIs(type(response),
messages_pb2.StreamingOutputCallResponse)
self.assertEqual(_RESPONSE_PAYLOAD_SIZE,
len(response.payload.body))
# After the RPC is finished, further reads will lead to exception.
self.assertEqual(await call.code(), grpc.StatusCode.OK)
with self.assertRaises(asyncio.InvalidStateError):
await call.read()
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_early_cancel_unary_stream(self):
"""Test cancellation before receiving messages."""
async with aio.insecure_channel(self._server_target) as channel:
stub = test_pb2_grpc.TestServiceStub(channel)
# Prepares the request
request = messages_pb2.StreamingOutputCallRequest()
for _ in range(_NUM_STREAM_RESPONSES):
request.response_parameters.append(
messages_pb2.ResponseParameters(
size=_RESPONSE_PAYLOAD_SIZE,
interval_us=_RESPONSE_INTERVAL_US,
))
# Invokes the actual RPC
call = stub.StreamingOutputCall(request)
self.assertFalse(call.cancelled())
self.assertTrue(call.cancel())
self.assertFalse(call.cancel())
with self.assertRaises(asyncio.CancelledError):
await call.read()
self.assertTrue(call.cancelled())
self.assertEqual(grpc.StatusCode.CANCELLED, await call.code())
self.assertEqual(_LOCAL_CANCEL_DETAILS_EXPECTATION, await
call.details())
def setUp(self):
self.server = test_common.test_server()
test_pb2_grpc.add_TestServiceServicer_to_server(
service.TestService(), self.server)
port = self.server.add_insecure_port('[::]:0')
self.server.start()
self.stub = test_pb2_grpc.TestServiceStub(
grpc.insecure_channel('localhost:{}'.format(port)))
def child_target(parent_bidi_call, parent_channel, args):
stub = test_pb2_grpc.TestServiceStub(parent_channel)
try:
_blocking_unary(stub)
raise Exception(
'Child should not be able to re-use channel after fork')
except ValueError as expected_value_error:
pass
def setUp(self):
self.server = grpc.server(futures.ThreadPoolExecutor(max_workers=10))
test_pb2_grpc.add_TestServiceServicer_to_server(
methods.TestService(), self.server)
port = self.server.add_insecure_port('[::]:0')
self.server.start()
self.stub = test_pb2_grpc.TestServiceStub(
grpc.insecure_channel('localhost:{}'.format(port)))
def _close_channel_before_fork(channel, args):
def child_target():
new_channel.close()
with _channel(args) as child_channel:
child_stub = test_pb2_grpc.TestServiceStub(child_channel)
_blocking_unary(child_stub)
stub = test_pb2_grpc.TestServiceStub(channel)
_blocking_unary(stub)
channel.close()
with _channel(args) as new_channel:
new_stub = test_pb2_grpc.TestServiceStub(new_channel)
child_process = _ChildProcess(child_target)
child_process.start()
_blocking_unary(new_stub)
child_process.finish()
async def ping_pong():
address, server = await start_test_server()
channel = aio.insecure_channel(address)
stub = test_pb2_grpc.TestServiceStub(channel)
await stub.UnaryCall(messages_pb2.SimpleRequest())
await channel.close()
await server.stop(None)
async def test_close_async_context(self):
async with aio.insecure_channel(self._server_target) as channel:
stub = test_pb2_grpc.TestServiceStub(channel)
calls = [
stub.UnaryCall(messages_pb2.SimpleRequest()) for _ in range(2)
]
for call in calls:
self.assertTrue(call.cancelled())
async def _create_server_stub_pair(
*interceptors: aio.ServerInterceptor
) -> Tuple[aio.Server, test_pb2_grpc.TestServiceStub]:
"""Creates a server-stub pair with given interceptors.
Returning the server object to protect it from being garbage collected.
"""
server_target, server = await start_test_server(interceptors=interceptors)
channel = aio.insecure_channel(server_target)
return server, test_pb2_grpc.TestServiceStub(channel)
async def async_work():
# Create async stub
async_channel = aio.insecure_channel(address,
options=_unique_options())
async_stub = test_pb2_grpc.TestServiceStub(async_channel)
call = async_stub.UnaryCall(messages_pb2.SimpleRequest())
response = await call
self.assertIsInstance(response, messages_pb2.SimpleResponse)
self.assertEqual(grpc.StatusCode.OK, await call.code())
