Add tests to ensure KernelManager subclass methods are called

jupyter_client/tests/test_kernelmanager.py

@@ -22,7 +22,7 @@
 
 from ..manager import start_new_kernel, start_new_async_kernel
 from ..manager import _ShutdownStatus
-from .utils import test_env, skip_win32, AsyncKernelManagerSubclass, AsyncKernelManagerWithCleanup
+from .utils import test_env, SyncKernelManagerSubclass, AsyncKernelManagerSubclass, AsyncKernelManagerWithCleanup
 
 pjoin = os.path.join
 
@@ -104,6 +104,11 @@ def km(config):
     km = KernelManager(config=config)
     return km
 
+@pytest.fixture
+def km_subclass(config):
+    km = SyncKernelManagerSubclass(config=config)
+    return km
+
 
 @pytest.fixture
 def zmq_context():
@@ -119,6 +124,12 @@ def async_km(request, config):
     return km
 
 
+@pytest.fixture
+def async_km_subclass(config):
+    km = AsyncKernelManagerSubclass(config=config)
+    return km
+
+
 @pytest.fixture
 @async_generator  # This is only necessary while Python 3.5 is support afterwhich both it and yield_() can be removed
 async def start_async_kernel():
@@ -283,6 +294,45 @@ def test_no_cleanup_shared_context(self, zmq_context):
         assert km.context.closed is False
         assert zmq_context.closed is False
 
+    def test_subclass_callables(self, km_subclass):
+        km_subclass.reset_counts()
+        km_subclass.start_kernel(stdout=PIPE, stderr=PIPE)
+        assert km_subclass.call_count('start_kernel') == 1
+        assert km_subclass.call_count('_launch_kernel') == 1
+
+        is_alive = km_subclass.is_alive()
+        assert is_alive
+        km_subclass.reset_counts()
+        km_subclass.restart_kernel(now=True)
+        assert km_subclass.call_count('restart_kernel') == 1
+        assert km_subclass.call_count('shutdown_kernel') == 1
+        assert km_subclass.call_count('interrupt_kernel') == 1
+        assert km_subclass.call_count('_kill_kernel') == 1
+        assert km_subclass.call_count('cleanup_resources') == 1
+        assert km_subclass.call_count('start_kernel') == 1
+        assert km_subclass.call_count('_launch_kernel') == 1
+
+        is_alive = km_subclass.is_alive()
+        assert is_alive
+
+        km_subclass.reset_counts()
+        km_subclass.interrupt_kernel()
+        assert km_subclass.call_count('interrupt_kernel') == 1
+
+        assert isinstance(km_subclass, KernelManager)
+
+        km_subclass.reset_counts()
+        km_subclass.shutdown_kernel(now=False)
+        assert km_subclass.call_count('shutdown_kernel') == 1
+        assert km_subclass.call_count('interrupt_kernel') == 1
+        assert km_subclass.call_count('request_shutdown') == 1
+        assert km_subclass.call_count('finish_shutdown') == 1
+        assert km_subclass.call_count('cleanup_resources') == 1
+
+        is_alive = km_subclass.is_alive()
+        assert is_alive is False
+        assert km_subclass.context.closed
+
 
 class TestParallel:
 
@@ -472,3 +522,42 @@ async def test_start_new_async_kernel(self, install_kernel, start_async_kernel):
         assert is_alive
         is_alive = await kc.is_alive()
         assert is_alive
+
+    async def test_subclass_callables(self, async_km_subclass):
+        async_km_subclass.reset_counts()
+        await async_km_subclass.start_kernel(stdout=PIPE, stderr=PIPE)
+        assert async_km_subclass.call_count('start_kernel') == 1
+        assert async_km_subclass.call_count('_launch_kernel') == 1
+
+        is_alive = await async_km_subclass.is_alive()
+        assert is_alive
+        async_km_subclass.reset_counts()
+        await async_km_subclass.restart_kernel(now=True)
+        assert async_km_subclass.call_count('restart_kernel') == 1
+        assert async_km_subclass.call_count('shutdown_kernel') == 1
+        assert async_km_subclass.call_count('interrupt_kernel') == 1
+        assert async_km_subclass.call_count('_kill_kernel') == 1
+        assert async_km_subclass.call_count('cleanup_resources') == 1
+        assert async_km_subclass.call_count('start_kernel') == 1
+        assert async_km_subclass.call_count('_launch_kernel') == 1
+
+        is_alive = await async_km_subclass.is_alive()
+        assert is_alive
+
+        async_km_subclass.reset_counts()
+        await async_km_subclass.interrupt_kernel()
+        assert async_km_subclass.call_count('interrupt_kernel') == 1
+
+        assert isinstance(async_km_subclass, AsyncKernelManager)
+
+        async_km_subclass.reset_counts()
+        await async_km_subclass.shutdown_kernel(now=False)
+        assert async_km_subclass.call_count('shutdown_kernel') == 1
+        assert async_km_subclass.call_count('interrupt_kernel') == 1
+        assert async_km_subclass.call_count('request_shutdown') == 1
+        assert async_km_subclass.call_count('finish_shutdown') == 1
+        assert async_km_subclass.call_count('cleanup_resources') == 1
+
+        is_alive = await async_km_subclass.is_alive()
+        assert is_alive is False
+        assert async_km_subclass.context.closed

jupyter_client/tests/utils.py

@@ -6,9 +6,10 @@
 import sys
 from unittest.mock import patch
 from tempfile import TemporaryDirectory
+from typing import Dict
 
 import pytest
-from jupyter_client import AsyncKernelManager
+from jupyter_client import AsyncKernelManager, KernelManager
 
 
 skip_win32 = pytest.mark.skipif(sys.platform.startswith('win'), reason="Windows")
@@ -61,22 +62,115 @@ def execute(code='', kc=None, **kwargs):
     return msg_id, reply['content']
 
 
-class AsyncKernelManagerSubclass(AsyncKernelManager):
-    """Used to test deprecation "routes" that are determined by superclass' detection of methods.
+class RecordCallMixin:
+    method_calls: Dict[str, int] = {}
+
+    def record(self, method_name: str) -> None:
+        if method_name not in self.method_calls:
+            self.method_calls[method_name] = 0
+        self.method_calls[method_name] += 1
+
+    def call_count(self, method_name: str) -> int:
+        assert method_name in self.method_calls
+        return self.method_calls[method_name]
+
+    def reset_counts(self) -> None:
+        for record in self.method_calls.keys():
+            self.method_calls[record] = 0
+
+
+class SyncKernelManagerSubclass(RecordCallMixin, KernelManager):
+
+    def start_kernel(self, **kw):
+        self.record('start_kernel')
+        return super().start_kernel(**kw)
+
+    def shutdown_kernel(self, now=False, restart=False):
+        self.record('shutdown_kernel')
+        return super().shutdown_kernel(now=now, restart=restart)
 
-       This class represents a current subclass that overrides both cleanup() and cleanup_resources()
-       in order to be compatible with older jupyter_clients.  We should find that cleanup_resources()
-       is called on these instances vix TestAsyncKernelManagerSubclass.
+    def restart_kernel(self, now=False, **kw):
+        self.record('restart_kernel')
+        return super().restart_kernel(now=now, **kw)
+
+    def interrupt_kernel(self):
+        self.record('interrupt_kernel')
+        return super().interrupt_kernel()
+
+    def request_shutdown(self, restart=False):
+        self.record('request_shutdown')
+        return super().request_shutdown(restart=restart)
+
+    def finish_shutdown(self, waittime=None, pollinterval=0.1):
+        self.record('finish_shutdown')
+        return super().finish_shutdown(waittime=waittime, pollinterval=pollinterval)
+
+    def _launch_kernel(self, kernel_cmd, **kw):
+        self.record('_launch_kernel')
+        return super()._launch_kernel(kernel_cmd, **kw)
+
+    def _kill_kernel(self):
+        self.record('_kill_kernel')
+        return super()._kill_kernel()
+
+    def cleanup_resources(self, restart=False):
+        self.record('cleanup_resources')
+        super().cleanup_resources(restart=restart)
+
+
+class AsyncKernelManagerSubclass(RecordCallMixin, AsyncKernelManager):
+    """Used to test subclass hierarchies to ensure methods are called when expected.
+
+       This class is also used to test deprecation "routes" that are determined by superclass'
+       detection of methods.
+
+       This class represents a current subclass that overrides "interesting" methods of AsyncKernelManager.
     """
+    which_cleanup = ""  # cleanup deprecation testing
+
+    async def start_kernel(self, **kw):
+        self.record('start_kernel')
+        return await super().start_kernel(**kw)
+
+    async def shutdown_kernel(self, now=False, restart=False):
+        self.record('shutdown_kernel')
+        return await super().shutdown_kernel(now=now, restart=restart)
+
+    async def restart_kernel(self, now=False, **kw):
+        self.record('restart_kernel')
+        return await super().restart_kernel(now=now, **kw)
+
+    async def interrupt_kernel(self):
+        self.record('interrupt_kernel')
+        return await super().interrupt_kernel()
+
+    def request_shutdown(self, restart=False):
+        self.record('request_shutdown')
+        return super().request_shutdown(restart=restart)
+
+    async def finish_shutdown(self, waittime=None, pollinterval=0.1):
+        self.record('finish_shutdown')
+        return await super().finish_shutdown(waittime=waittime, pollinterval=pollinterval)
+
+    async def _launch_kernel(self, kernel_cmd, **kw):
+        self.record('_launch_kernel')
+        return await super()._launch_kernel(kernel_cmd, **kw)
+
+    async def _kill_kernel(self):
+        self.record('_kill_kernel')
+        return await super()._kill_kernel()
 
     def cleanup(self, connection_file=True):
+        self.record('cleanup')
         super().cleanup(connection_file=connection_file)
         self.which_cleanup = 'cleanup'
 
     def cleanup_resources(self, restart=False):
+        self.record('cleanup_resources')
         super().cleanup_resources(restart=restart)
         self.which_cleanup = 'cleanup_resources'
 
+
 class AsyncKernelManagerWithCleanup(AsyncKernelManager):
     """Used to test deprecation "routes" that are determined by superclass' detection of methods.