From c5c811d4ce761ec9ea6d78220d6f49605e7dd2f7 Mon Sep 17 00:00:00 2001
From: hanhxiao <hanhxiao@tencent.com>
Date: Fri, 2 Aug 2019 13:03:27 +0800
Subject: [PATCH] fix(test): fix grpc gentle shutdown

---
 gnes/helper.py            | 2 +-
 gnes/service/base.py      | 2 +-
 gnes/service/grpc.py      | 2 +-
 tests/test_stream_grpc.py | 2 +-
 4 files changed, 4 insertions(+), 4 deletions(-)

diff --git a/gnes/helper.py b/gnes/helper.py
index 7e69b0bf..f732ca68 100644
--- a/gnes/helper.py
+++ b/gnes/helper.py
@@ -194,6 +194,7 @@ class ColoredFormatter(Formatter):
         'WARNING': dict(color='red', on_color='on_yellow'),  # yellow
         'ERROR': dict(color='white', on_color='on_red'),  # 31 for red
         'CRITICAL': dict(color='red', on_color='on_white'),  # white on red bg
+        'SUCCESS': dict(color='white', on_color='on_green'),  # green
     }
 
     PREFIX = '\033['
@@ -535,4 +536,3 @@ def load_contrib_module():
 profiling = time_profile
 
 yaml = _get_yaml()
-
diff --git a/gnes/service/base.py b/gnes/service/base.py
index bfee42d2..be9972cc 100644
--- a/gnes/service/base.py
+++ b/gnes/service/base.py
@@ -264,7 +264,7 @@ def _run(self, ctx):
             self.is_ready.set()
             self.is_event_loop.set()
             self._start_auto_dump()
-            self.logger.info('ready and listening')
+            self.logger.info(colored('ready and listening', color='white', on_color='on_green'))
             while self.is_event_loop.is_set():
                 pull_sock = None
                 socks = dict(poller.poll())
diff --git a/gnes/service/grpc.py b/gnes/service/grpc.py
index 231f8e17..f89fdf22 100644
--- a/gnes/service/grpc.py
+++ b/gnes/service/grpc.py
@@ -161,4 +161,4 @@ def __enter__(self):
         return self
 
     def __exit__(self, exc_type, exc_val, exc_tb):
-        self.server.stop()
+        self.server.stop(None)
diff --git a/tests/test_stream_grpc.py b/tests/test_stream_grpc.py
index ad80ae1d..f500fd20 100644
--- a/tests/test_stream_grpc.py
+++ b/tests/test_stream_grpc.py
@@ -81,7 +81,7 @@ def test_async_block(self):
             '--socket_out', str(SocketType.PUSH_CONNECT),
         ])
 
-        with Router1(p1_args), Router2(p2_args), GRPCFrontend(args), grpc.insecure_channel(
+        with GRPCFrontend(args), Router1(p1_args), Router2(p2_args), grpc.insecure_channel(
                 '%s:%s' % (args.grpc_host, args.grpc_port),
                 options=[('grpc.max_send_message_length', 70 * 1024 * 1024),
                          ('grpc.max_receive_message_length', 70 * 1024 * 1024)]) as channel: