all loop to ensure kernel is alive before connecting working unit tests

Author: Zsailer

jupyter_server/services/kernels/connection/channels.py

@@ -1,5 +1,6 @@
 import asyncio
 import json
+import time
 import weakref
 from concurrent.futures import Future
 from textwrap import dedent
@@ -16,6 +17,8 @@
 except ImportError:
     from jupyter_client.jsonutil import date_default as json_default
 
+from jupyter_client.utils import ensure_async
+
 from jupyter_server.transutils import _i18n
 
 from .abc import KernelWebsocketConnectionABC
@@ -282,7 +285,7 @@ async def _register_session(self):
         if (
             self.kernel_id in self.multi_kernel_manager
         ):  # only update open sessions if kernel is actively managed
-            self._open_sessions[self.session_key] = self
+            self._open_sessions[self.session_key] = self.websocket_handler
 
     async def prepare(self):
         # check session collision:
@@ -302,6 +305,12 @@ async def prepare(self):
                 self.kernel_manager.reason = str(e)
                 raise web.HTTPError(500, str(e)) from e
 
+        t0 = time.time()
+        while not await ensure_async(self.kernel_manager.is_alive()):
+            await asyncio.sleep(0.1)
+            if time.time() - t0 > self.multi_kernel_manager.kernel_info_timeout:
+                raise TimeoutError("Kernel never reached an 'alive' state.")
+
         self.session.key = self.kernel_manager.session.key
         future = self.request_kernel_info()
 
@@ -360,7 +369,7 @@ def replay(value):
                 for _, stream in self.channels.items():
                     if not stream.closed():
                         stream.close()
-                self.close()
+                self.disconnect()
                 return
 
         self.multi_kernel_manager.add_restart_callback(self.kernel_id, self.on_kernel_restarted)
@@ -376,6 +385,9 @@ def subscribe(value):
         ZMQChannelsWebsocketConnection._open_sockets.add(self)
         return connected
 
+    def close(self):
+        return self.disconnect()
+
     def disconnect(self):
         self.log.debug("Websocket closed %s", self.session_key)
         # unregister myself as an open session (only if it's really me)
@@ -536,7 +548,7 @@ def _on_zmq_reply(self, stream, msg_list):
         # eventloop but hasn't been called.
         if stream.closed():
             self.log.warning("zmq message arrived on closed channel")
-            self.close()
+            self.disconnect()
             return
         channel = getattr(stream, "channel", None)
         if self.subprotocol == "v1.kernel.websocket.jupyter.org":

jupyter_server/services/kernels/websocket.py

@@ -178,8 +178,7 @@ async def pre_get(self):
 
         kernel = self.kernel_manager.get_kernel(self.kernel_id)
         self.connection = self.kernel_websocket_connection_class(
-            parent=kernel,
-            websocket_handler=self,
+            parent=kernel, websocket_handler=self, config=self.config
         )
 
         if self.get_argument("session_id", None):

tests/services/sessions/test_api.py

@@ -20,7 +20,7 @@
 
 from ...utils import expected_http_error
 
-TEST_TIMEOUT = 60
+TEST_TIMEOUT = 10
 
 
 @pytest.fixture(autouse=True)
@@ -556,9 +556,13 @@ async def test_restart_kernel(session_client, jp_base_url, jp_fetch, jp_ws_fetch
     model = json.loads(r.body.decode())
     assert model["connections"] == 0
 
-    # Open a websocket connection.
-    await jp_ws_fetch("api", "kernels", kid, "channels")
+    # Add a delay to give the kernel enough time to restart.
+    #    time.sleep(2)
 
+    # Open a websocket connection.
+    ws2 = await jp_ws_fetch("api", "kernels", kid, "channels")
+    # Close/open websocket
+    ws2.close()
     r = await jp_fetch("api", "kernels", kid, method="GET")
     model = json.loads(r.body.decode())
     assert model["connections"] == 1