work updates to tests

Author: Zsailer

jupyter_client/manager.py

@@ -89,6 +89,7 @@ class KernelManager(ConnectionFileMixin):
     def __init__(self, *args, **kwargs):
         super().__init__(**kwargs)
         self._shutdown_status = _ShutdownStatus.Unset
+        # Create a place holder future.
         try:
             self._ready = Future()
         except RuntimeError:

jupyter_client/multikernelmanager.py

@@ -97,7 +97,12 @@ def create_kernel_manager(*args, **kwargs) -> KernelManager:
 
     context = Instance("zmq.Context")
 
-    _starting_kernels = Dict()
+    _pending_kernels = Dict()
+
+    @property
+    def _starting_kernels(self):
+        """A shim for backwards compatibility."""
+        return self._pending_kernels
 
     @default("context")
     def _context_default(self) -> zmq.Context:
@@ -165,7 +170,22 @@ async def _add_kernel_when_ready(
             await kernel_awaitable
             self._kernels[kernel_id] = km
         finally:
-            self._starting_kernels.pop(kernel_id, None)
+            self._pending_kernels.pop(kernel_id, None)
+
+    async def _remove_kernel_when_ready(
+        self, kernel_id: str, kernel_awaitable: t.Awaitable
+    ) -> None:
+        try:
+            await kernel_awaitable
+            self.remove_kernel(kernel_id)
+        finally:
+            self._pending_kernels.pop(kernel_id, None)
+
+    def _using_pending_kernels(self):
+        """Returns a boolean; a clearer method for determining if
+        this multikernelmanager is using pending kernels or not
+        """
+        return getattr(self, 'use_pending_kernels', False)
 
     async def _async_start_kernel(self, kernel_name: t.Optional[str] = None, **kwargs) -> str:
         """Start a new kernel.
@@ -186,12 +206,18 @@ async def _async_start_kernel(self, kernel_name: t.Optional[str] = None, **kwarg
 
         starter = ensure_async(km.start_kernel(**kwargs))
         fut = asyncio.ensure_future(self._add_kernel_when_ready(kernel_id, km, starter))
-        self._starting_kernels[kernel_id] = fut
+        self._pending_kernels[kernel_id] = fut
 
-        if getattr(self, 'use_pending_kernels', False):
+        # Handling a Pending Kernel
+        if self._using_pending_kernels():
+            # If using pending kernels, do not block
+            # on the kernel start.
             self._kernels[kernel_id] = km
         else:
             await fut
+            # raise an exception if one occurred during kernel startup.
+            if km.ready.exception():
+                raise km.ready.exception()
 
         return kernel_id
 
@@ -215,23 +241,28 @@ async def _async_shutdown_kernel(
             Will the kernel be restarted?
         """
         kernel = self.get_kernel(kernel_id)
-        if getattr(self, 'use_pending_kernels', False):
+        if self._using_pending_kernels():
             # Make sure the previous kernel started before trying to shutdown.
             if not kernel.ready.done():
                 raise RuntimeError("Kernel is in a pending state. Cannot shutdown.")
             # If the kernel didn't start properly, no need to shutdown.
             elif kernel.ready.exception():
+                self.remove_kernel(kernel_id)
                 return
         self.log.info("Kernel shutdown: %s" % kernel_id)
-        if kernel_id in self._starting_kernels:
+        if kernel_id in self._pending_kernels:
             try:
-                await self._starting_kernels[kernel_id]
+                await self._pending_kernels[kernel_id]
             except Exception:
                 self.remove_kernel(kernel_id)
                 return
         km = self.get_kernel(kernel_id)
-        await ensure_async(km.shutdown_kernel(now, restart))
-        self.remove_kernel(kernel_id)
+        stopper = ensure_async(km.shutdown_kernel(now, restart))
+        fut = asyncio.ensure_future(self._remove_kernel_when_ready(kernel_id, stopper))
+        self._pending_kernels[kernel_id] = fut
+        # Await the kernel if not using pending kernels.
+        if not self._using_pending_kernels():
+            await fut
 
     shutdown_kernel = run_sync(_async_shutdown_kernel)
 
@@ -266,13 +297,13 @@ def remove_kernel(self, kernel_id: str) -> KernelManager:
     async def _async_shutdown_all(self, now: bool = False) -> None:
         """Shutdown all kernels."""
         kids = self.list_kernel_ids()
-        kids += list(self._starting_kernels)
+        kids += list(self._pending_kernels)
         futs = [ensure_async(self.shutdown_kernel(kid, now=now)) for kid in set(kids)]
         await asyncio.gather(*futs)
 
     shutdown_all = run_sync(_async_shutdown_all)
 
-    @kernel_method
+    # @kernel_method
     def interrupt_kernel(self, kernel_id: str) -> None:
         """Interrupt (SIGINT) the kernel by its uuid.
 
@@ -281,7 +312,12 @@ def interrupt_kernel(self, kernel_id: str) -> None:
         kernel_id : uuid
             The id of the kernel to interrupt.
         """
+        kernel = self.get_kernel(kernel_id)
+        if not kernel.ready.done():
+            raise RuntimeError("Kernel is in a pending state. Cannot interrupt.")
+        out = kernel.interrupt_kernel()
         self.log.info("Kernel interrupted: %s" % kernel_id)
+        return out
 
     @kernel_method
     def signal_kernel(self, kernel_id: str, signum: int) -> None:
@@ -315,7 +351,7 @@ async def _async_restart_kernel(self, kernel_id: str, now: bool = False) -> None
             it is given a chance to perform a clean shutdown or not.
         """
         kernel = self.get_kernel(kernel_id)
-        if getattr(self, 'use_pending_kernels', False):
+        if self._using_pending_kernels():
             if not kernel.ready.done():
                 raise RuntimeError("Kernel is in a pending state. Cannot restart.")
         out = await ensure_async(kernel.restart_kernel(now=now))

jupyter_client/tests/test_multikernelmanager.py

@@ -4,6 +4,7 @@
 import os
 import sys
 import uuid
+from asyncio import ensure_future
 from subprocess import PIPE
 from unittest import TestCase
 
@@ -354,7 +355,7 @@ async def test_shutdown_all_while_starting(self):
         self.assertNotIn(kid, km)
 
         # Start another kernel
-        kid = await km.start_kernel(stdout=PIPE, stderr=PIPE)
+        kid = await ensure_future(km.start_kernel(stdout=PIPE, stderr=PIPE))
         self.assertIn(kid, km)
         self.assertEqual(len(km), 1)
         await km.shutdown_all()
@@ -365,52 +366,69 @@ async def test_shutdown_all_while_starting(self):
     @gen_test
     async def test_use_pending_kernels(self):
         km = self._get_pending_kernels_km()
-        kid = await now(km.start_kernel(stdout=PIPE, stderr=PIPE))
+        kid = await ensure_future(km.start_kernel(stdout=PIPE, stderr=PIPE))
         kernel = km.get_kernel(kid)
         assert not kernel.ready.done()
         assert kid in km
         assert kid in km.list_kernel_ids()
         assert len(km) == 1, f"{len(km)} != {1}"
-        await now(kernel.ready)
-        await now(km.restart_kernel(kid, now=True))
-        assert await km.is_alive(kid)
+        # Wait for the kernel to start.
+        await kernel.ready
+        await km.restart_kernel(kid, now=True)
+        out = await km.is_alive(kid)
+        assert out
         assert kid in km.list_kernel_ids()
-        await now(km.interrupt_kernel(kid))
+        await km.interrupt_kernel(kid)
         k = km.get_kernel(kid)
         assert isinstance(k, AsyncKernelManager)
-        await now(km.shutdown_kernel(kid, now=True))
+        await ensure_future(km.shutdown_kernel(kid, now=True))
+        # Wait for the kernel to shutdown
+        await kernel.ready
         assert kid not in km, f"{kid} not in {km}"
 
     @gen_test
     async def test_use_pending_kernels_early_restart(self):
         km = self._get_pending_kernels_km()
-        kid = await now(km.start_kernel(stdout=PIPE, stderr=PIPE))
+        kid = await ensure_future(km.start_kernel(stdout=PIPE, stderr=PIPE))
         kernel = km.get_kernel(kid)
         assert not kernel.ready.done()
         with pytest.raises(RuntimeError):
             await km.restart_kernel(kid, now=True)
-        await now(kernel.ready)
-        await now(km.shutdown_kernel(kid, now=True))
+        await kernel.ready
+        await ensure_future(km.shutdown_kernel(kid, now=True))
+        # Wait for the kernel to shutdown
+        await kernel.ready
         assert kid not in km, f"{kid} not in {km}"
 
     @gen_test
     async def test_use_pending_kernels_early_shutdown(self):
         km = self._get_pending_kernels_km()
-        kid = await km.start_kernel(stdout=PIPE, stderr=PIPE)
+        kid = await ensure_future(km.start_kernel(stdout=PIPE, stderr=PIPE))
         kernel = km.get_kernel(kid)
         assert not kernel.ready.done()
-        await km.shutdown_kernel(kid, now=True)
+        # Try shutting down while the kernel is pending
+        with pytest.raises(RuntimeError):
+            await ensure_future(km.shutdown_kernel(kid, now=True))
+        await kernel.ready
+        # Shutdown once the kernel is ready
+        await ensure_future(km.shutdown_kernel(kid, now=True))
+        # Wait for the kernel to shutdown
+        await kernel.ready
         assert kid not in km, f"{kid} not in {km}"
 
     @gen_test
     async def test_use_pending_kernels_early_interrupt(self):
         km = self._get_pending_kernels_km()
-        kid = await km.start_kernel(stdout=PIPE, stderr=PIPE)
+        kid = await ensure_future(km.start_kernel(stdout=PIPE, stderr=PIPE))
         kernel = km.get_kernel(kid)
         assert not kernel.ready.done()
         with pytest.raises(RuntimeError):
             await km.interrupt_kernel(kid)
-        await km.shutdown_kernel(kid, now=True)
+        # Now wait for the kernel to be ready.
+        await kernel.ready
+        await ensure_future(km.shutdown_kernel(kid, now=True))
+        # Wait for the kernel to shutdown
+        await kernel.ready
         assert kid not in km, f"{kid} not in {km}"
 
     @gen_test
@@ -555,7 +573,7 @@ async def test_bad_kernelspec(self):
             name="bad",
         )
         with pytest.raises(FileNotFoundError):
-            await km.start_kernel(kernel_name="bad", stdout=PIPE, stderr=PIPE)
+            await ensure_future(km.start_kernel(kernel_name="bad", stdout=PIPE, stderr=PIPE))
 
     @gen_test
     async def test_bad_kernelspec_pending(self):
@@ -565,10 +583,11 @@ async def test_bad_kernelspec_pending(self):
             argv=["non_existent_executable"],
             name="bad",
         )
-        kernel_id = await km.start_kernel(kernel_name="bad", stdout=PIPE, stderr=PIPE)
-        assert kernel_id in km._starting_kernels
+        kernel_id = await ensure_future(
+            km.start_kernel(kernel_name="bad", stdout=PIPE, stderr=PIPE)
+        )
         with pytest.raises(FileNotFoundError):
             await km.get_kernel(kernel_id).ready
         assert kernel_id in km.list_kernel_ids()
-        await km.shutdown_kernel(kernel_id)
+        await ensure_future(km.shutdown_kernel(kernel_id))
         assert kernel_id not in km.list_kernel_ids()