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())