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_unary_unary(self):
call = self._stub.UnaryCall(messages_pb2.SimpleRequest())
validation = inject_callbacks(call)
self.assertEqual(grpc.StatusCode.OK, await call.code())
await validation
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_handler(request: bytes, context: aio.ServicerContext):
self.assertEqual(_REQUEST, request)
def exception_raiser(unused_context):
raise RuntimeError('A test RuntimeError')
context.add_done_callback(exception_raiser)
validation_future.set_result(inject_callbacks(context))
return _RESPONSE
async def test_stream_unary(self):
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
call = self._stub.StreamingInputCall(gen())
validation = inject_callbacks(call)
response = await call
self.assertIsInstance(response, messages_pb2.StreamingInputCallResponse)
self.assertEqual(_NUM_STREAM_RESPONSES * _REQUEST_PAYLOAD_SIZE,
response.aggregated_payload_size)
self.assertEqual(grpc.StatusCode.OK, await call.code())
await validation
async def test_stream_stream(self):
call = self._stub.FullDuplexCall()
validation = inject_callbacks(call)
request = messages_pb2.StreamingOutputCallRequest()
request.response_parameters.append(
messages_pb2.ResponseParameters(size=_RESPONSE_PAYLOAD_SIZE))
for _ in range(_NUM_STREAM_RESPONSES):
await call.write(request)
response = await call.read()
self.assertIsInstance(response,
messages_pb2.StreamingOutputCallResponse)
self.assertEqual(_RESPONSE_PAYLOAD_SIZE, len(response.payload.body))
await call.done_writing()
self.assertEqual(grpc.StatusCode.OK, await call.code())
await validation
async def test_unary_stream(self):
request = messages_pb2.StreamingOutputCallRequest()
for _ in range(_NUM_STREAM_RESPONSES):
request.response_parameters.append(
messages_pb2.ResponseParameters(size=_RESPONSE_PAYLOAD_SIZE))
call = self._stub.StreamingOutputCall(request)
validation = inject_callbacks(call)
response_cnt = 0
async for response in call:
response_cnt += 1
self.assertIsInstance(response,
messages_pb2.StreamingOutputCallResponse)
self.assertEqual(_RESPONSE_PAYLOAD_SIZE, len(response.payload.body))
self.assertEqual(_NUM_STREAM_RESPONSES, response_cnt)
self.assertEqual(grpc.StatusCode.OK, await call.code())
await validation
async def test_handler(request: bytes, context: aio.ServicerContext):
self.assertEqual(_REQUEST, request)
validation_future.set_result(inject_callbacks(context))
raise RuntimeError('A test RuntimeError')
async def test_handler(request_iterator, context: aio.ServicerContext):
validation_future.set_result(inject_callbacks(context))
async for request in request_iterator:
self.assertEqual(_REQUEST, request)
return _RESPONSE
async def test_handler(request: bytes, context: aio.ServicerContext):
self.assertEqual(_REQUEST, request)
validation_future.set_result(inject_callbacks(context))
for _ in range(_NUM_STREAM_RESPONSES):
yield _RESPONSE
async def test_handler(request: bytes, context: aio.ServicerContext):
self.assertEqual(_REQUEST, request)
validation_future.set_result(inject_callbacks(context))
return _RESPONSE
