Add test case for set_code with no return value

src/python/grpcio_tests/tests_aio/unit/server_test.py

@@ -38,6 +38,8 @@ _STREAM_STREAM_READER_WRITER = '/test/StreamStreamReaderWriter'
 _STREAM_STREAM_EVILLY_MIXED = '/test/StreamStreamEvillyMixed'
 _UNIMPLEMENTED_METHOD = '/test/UnimplementedMethod'
 _ERROR_IN_STREAM_STREAM = '/test/ErrorInStreamStream'
+_ERROR_WITHOUT_RAISE_IN_UNARY_UNARY = '/test/ErrorWithoutRaiseInUnaryUnary'
+_ERROR_WITHOUT_RAISE_IN_STREAM_STREAM = '/test/ErrorWithoutRaiseInStreamStream'
 
 _REQUEST = b'\x00\x00\x00'
 _RESPONSE = b'\x01\x01\x01'
@@ -86,6 +88,12 @@ class _GenericHandler(grpc.GenericRpcHandler):
             _ERROR_IN_STREAM_STREAM:
                 grpc.stream_stream_rpc_method_handler(
                     self._error_in_stream_stream),
+            _ERROR_WITHOUT_RAISE_IN_UNARY_UNARY:
+                grpc.unary_unary_rpc_method_handler(
+                    self._error_without_raise_in_unary_unary),
+            _ERROR_WITHOUT_RAISE_IN_STREAM_STREAM:
+                grpc.stream_stream_rpc_method_handler(
+                    self._error_without_raise_in_stream_stream),
         }
 
     @staticmethod
@@ -168,6 +176,16 @@ class _GenericHandler(grpc.GenericRpcHandler):
             raise RuntimeError('A testing RuntimeError!')
         yield _RESPONSE
 
+    async def _error_without_raise_in_unary_unary(self, request, context):
+        assert _REQUEST == request
+        context.set_code(grpc.StatusCode.INTERNAL)
+
+    async def _error_without_raise_in_stream_stream(self, request_iterator,
+                                                    context):
+        async for request in request_iterator:
+            assert _REQUEST == request
+        context.set_code(grpc.StatusCode.INTERNAL)
+
     def service(self, handler_details):
         self._called.set_result(None)
         return self._routing_table.get(handler_details.method)
@@ -426,6 +444,26 @@ class TestServer(AioTestBase):
         # Don't segfault here
         self.assertEqual(grpc.StatusCode.UNKNOWN, await call.code())
 
+    async def test_error_without_raise_in_unary_unary(self):
+        call = self._channel.unary_unary(_ERROR_WITHOUT_RAISE_IN_UNARY_UNARY)(
+            _REQUEST)
+
+        with self.assertRaises(aio.AioRpcError) as exception_context:
+            await call
+
+        rpc_error = exception_context.exception
+        self.assertEqual(grpc.StatusCode.INTERNAL, rpc_error.code())
+
+    async def test_error_without_raise_in_stream_stream(self):
+        call = self._channel.stream_stream(
+            _ERROR_WITHOUT_RAISE_IN_STREAM_STREAM)()
+
+        for _ in range(_NUM_STREAM_REQUESTS):
+            await call.write(_REQUEST)
+        await call.done_writing()
+
+        self.assertEqual(grpc.StatusCode.INTERNAL, await call.code())
+
 
 if __name__ == '__main__':
     logging.basicConfig(level=logging.DEBUG)