Add test for kernel client reset

Author: davidbrochart

nbclient/client.py

@@ -298,13 +298,11 @@ def __init__(self, nb, km=None, **kw):
         super().__init__(**kw)
         self.nb = nb
         self.km = km
-        self.reset_execution_trackers(reset_kc=True)
+        self.reset_execution_trackers()
 
-    def reset_execution_trackers(self, reset_kc=False):
+    def reset_execution_trackers(self):
         """Resets any per-execution trackers.
         """
-        if reset_kc:
-            self.kc = None
         self.code_cells_executed = 0
         self._display_id_map = {}
         self.widget_state = {}
@@ -418,6 +416,7 @@ async def async_setup_kernel(self, **kwargs):
         When control returns from the yield it stops the client's zmq channels, and shuts
         down the kernel.
         """
+        reset_kc = kwargs.pop('reset_kc', False)
         if self.km is None:
             self.start_kernel_manager()
 
@@ -426,7 +425,8 @@ async def async_setup_kernel(self, **kwargs):
         try:
             yield
         finally:
-            await self._async_cleanup_kernel()
+            if reset_kc:
+                await self._async_cleanup_kernel()
 
     async def async_execute(self, **kwargs):
         """
@@ -438,17 +438,19 @@ async def async_execute(self, **kwargs):
             Any option for `self.kernel_manager_class.start_kernel()`. Because
             that defaults to AsyncKernelManager, this will likely include options
             accepted by `AsyncKernelManager.start_kernel()``, which includes `cwd`.
-            If present, `reset_kc` is passed to `self.reset_execution_trackers`:
+            If present, `reset_kc` is passed to `self.async_setup_kernel`:
             if True, the kernel client will be reset and a new one will be created
-            (default: False).
+            and cleaned up after execution (default: False).
 
         Returns
         -------
         nb : NotebookNode
             The executed notebook.
         """
-        reset_kc = kwargs.pop('reset_kc', False)
-        self.reset_execution_trackers(reset_kc=reset_kc)
+        reset_kc = kwargs.get('reset_kc', False)
+        if reset_kc:
+            await self._async_cleanup_kernel()
+        self.reset_execution_trackers()
 
         async with self.async_setup_kernel(**kwargs):
             self.log.info("Executing notebook with kernel: %s" % self.kernel_name)

nbclient/tests/test_client.py

@@ -544,6 +544,29 @@ def test_force_raise_errors(self):
             else:
                 assert u"# üñîçø∂é".encode('utf8', 'replace') in str(exc.value)
 
+    def test_reset_kernel_client(self):
+        filename = os.path.join(current_dir, 'files', 'HelloWorld.ipynb')
+
+        with io.open(filename) as f:
+            input_nb = nbformat.read(f, 4)
+
+        executor = NotebookClient(
+            input_nb,
+            resources=self.build_resources(),
+        )
+
+        executor.execute()
+        # we didn't ask to reset the kernel client, a new one must have been created
+        kc = executor.kc
+        self.assertNotEqual(kc, None)
+        executor.execute()
+        # we didn't ask to reset the kernel client, the previously created one must have been reused
+        self.assertEqual(kc, executor.kc)
+        executor.execute(reset_kc=True)
+        # we asked to reset the kernel client, the previous one must have been cleaned up,
+        # a new one must have been created and also cleaned up
+        self.assertEqual(executor.kc, None)
+
     def test_custom_kernel_manager(self):
         from .fake_kernelmanager import FakeCustomKernelManager