Add cancel_all_calls to Python server

Also format _low_test.py to fit within the 80 character fill-limit and
re-style test assertions.

3 files changed, 21 insertions, 0 deletions

diff --git a/src/python/grpcio/grpc/_adapter/_c/types.h b/src/python/grpcio/grpc/_adapter/_c/types.h
index f646465c63..f6ff957baa 100644
--- a/src/python/grpcio/grpc/_adapter/_c/types.h
+++ b/src/python/grpcio/grpc/_adapter/_c/types.h
@@ -146,6 +146,7 @@ typedef struct Server {
   PyObject_HEAD
   grpc_server *c_serv;
   CompletionQueue *cq;
+  int shutdown_called;
 } Server;
 Server *pygrpc_Server_new(PyTypeObject *type, PyObject *args, PyObject *kwargs);
 void pygrpc_Server_dealloc(Server *self);
@@ -156,6 +157,7 @@ PyObject *pygrpc_Server_add_http2_port(
 PyObject *pygrpc_Server_start(Server *self, PyObject *ignored);
 PyObject *pygrpc_Server_shutdown(
     Server *self, PyObject *args, PyObject *kwargs);
+PyObject *pygrpc_Server_cancel_all_calls(Server *self, PyObject *unused);
 extern PyTypeObject pygrpc_Server_type;
 
 /*=========*/
diff --git a/src/python/grpcio/grpc/_adapter/_c/types/server.c b/src/python/grpcio/grpc/_adapter/_c/types/server.c
index 15c98f28eb..8feab8aab1 100644
--- a/src/python/grpcio/grpc/_adapter/_c/types/server.c
+++ b/src/python/grpcio/grpc/_adapter/_c/types/server.c
@@ -45,6 +45,8 @@ PyMethodDef pygrpc_Server_methods[] = {
      METH_KEYWORDS, ""},
     {"start", (PyCFunction)pygrpc_Server_start, METH_NOARGS, ""},
     {"shutdown", (PyCFunction)pygrpc_Server_shutdown, METH_KEYWORDS, ""},
+    {"cancel_all_calls", (PyCFunction)pygrpc_Server_cancel_all_calls,
+     METH_NOARGS, ""},
     {NULL}
 };
 const char pygrpc_Server_doc[] = "See grpc._adapter._types.Server.";
@@ -109,6 +111,7 @@ Server *pygrpc_Server_new(PyTypeObject *type, PyObject *args, PyObject *kwargs)
   pygrpc_discard_channel_args(c_args);
   self->cq = cq;
   Py_INCREF(self->cq);
+  self->shutdown_called = 0;
   return self;
 }
 
@@ -163,6 +166,7 @@ PyObject *pygrpc_Server_add_http2_port(
 
 PyObject *pygrpc_Server_start(Server *self, PyObject *ignored) {
   grpc_server_start(self->c_serv);
+  self->shutdown_called = 0;
   Py_RETURN_NONE;
 }
 
@@ -176,5 +180,17 @@ PyObject *pygrpc_Server_shutdown(
   }
   tag = pygrpc_produce_server_shutdown_tag(user_tag);
   grpc_server_shutdown_and_notify(self->c_serv, self->cq->c_cq, tag);
+  self->shutdown_called = 1;
+  Py_RETURN_NONE;
+}
+
+PyObject *pygrpc_Server_cancel_all_calls(Server *self, PyObject *unused) {
+  if (!self->shutdown_called) {
+    PyErr_SetString(
+        PyExc_RuntimeError,
+        "shutdown must have been called prior to calling cancel_all_calls!");
+    return NULL;
+  }
+  grpc_server_cancel_all_calls(self->c_serv);
   Py_RETURN_NONE;
 }
diff --git a/src/python/grpcio/grpc/_adapter/_low.py b/src/python/grpcio/grpc/_adapter/_low.py
index 147086e725..3859ebb0e2 100644
--- a/src/python/grpcio/grpc/_adapter/_low.py
+++ b/src/python/grpcio/grpc/_adapter/_low.py
@@ -124,3 +124,6 @@ class Server(_types.Server):
 
   def request_call(self, completion_queue, tag):
     return self.server.request_call(completion_queue.completion_queue, tag)
+
+  def cancel_all_calls(self):
+    return self.server.cancel_all_calls()