similify console start end (#175)

* similify console start end

* better test case

aiomonitor/monitor.py

@@ -103,13 +103,8 @@ def start(self) -> None:
         assert not self._started
 
         self._started = True
-        h, p = self._host, self._console_port
         self._event_loop_thread_id = threading.get_ident()
         self._ui_thread.start()
-        if self._console_enabled:
-            log.info('Starting console at %s:%d', h, p)
-            self._console_future = init_console_server(
-                self._host, self._console_port, self._locals, self._loop)
 
     @property
     def closed(self) -> bool:
@@ -132,15 +127,6 @@ def close(self) -> None:
         if not self._closed:
             self._closing.set()
             self._ui_thread.join()
-            if self._console_future is not None:
-                server = self._console_future.result(timeout=15)
-                coro = close_server(server)
-                if self._loop.is_running():
-                    fut = asyncio.run_coroutine_threadsafe(
-                        coro, loop=self._loop)
-                    fut.result(timeout=15)
-                else:
-                    self._loop.run_until_complete(coro)
             self._closed = True
 
     def _server(self) -> None:
@@ -406,10 +392,18 @@ def do_console(self) -> None:
             self._sout.flush()
             return
 
-        if self._console_future is not None:
-            self._console_future.result(timeout=3)
-
-        console_proxy(self._sin, self._sout, self._host, self._console_port)
+        h, p = self._host, self._console_port
+        log.info('Starting console at %s:%d', h, p)
+        fut = init_console_server(
+            self._host, self._console_port, self._locals, self._loop)
+        server = fut.result(timeout=3)
+        try:
+            console_proxy(
+                self._sin, self._sout, self._host, self._console_port)
+        finally:
+            coro = close_server(server)
+            close_fut = asyncio.run_coroutine_threadsafe(coro, loop=self._loop)
+            close_fut.result(timeout=15)
 
 
 def start_monitor(loop: Loop, *,

tests/test_monitor.py

@@ -58,16 +58,21 @@ def tn_client(monitor):
     tn.close()
 
 
-def test_ctor(loop, unused_port):
+@pytest.fixture(params=[True, False], ids=['console:True', 'console:False'])
+def console_enabled(request):
+    return request.param
 
-    with Monitor(loop, console_enabled=False):
+
+def test_ctor(loop, unused_port, console_enabled):
+
+    with Monitor(loop, console_enabled=console_enabled):
         loop.run_until_complete(asyncio.sleep(0.01, loop=loop))
 
-    with start_monitor(loop, console_enabled=False) as m:
+    with start_monitor(loop, console_enabled=console_enabled) as m:
         loop.run_until_complete(asyncio.sleep(0.01, loop=loop))
     assert m.closed
 
-    m = Monitor(loop, console_enabled=False)
+    m = Monitor(loop, console_enabled=console_enabled)
     m.start()
     try:
         loop.run_until_complete(asyncio.sleep(0.01, loop=loop))
@@ -76,12 +81,18 @@ def test_ctor(loop, unused_port):
         m.close()  # make sure call is idempotent
     assert m.closed
 
-    m = Monitor(loop, console_enabled=False)
+    m = Monitor(loop, console_enabled=console_enabled)
     m.start()
     with m:
         loop.run_until_complete(asyncio.sleep(0.01, loop=loop))
     assert m.closed
 
+    # make sure that monitor inside async func can exit correctly
+    async def f():
+        with Monitor(loop, console_enabled=console_enabled):
+            await asyncio.sleep(0.01, loop=loop)
+    loop.run_until_complete(f())
+
 
 def execute(tn, command, pattern=b'>>>'):
     tn.write(command.encode('utf-8'))