reuse more kernels in kernel tests

IPython/kernel/multikernelmanager.py

@@ -171,7 +171,7 @@ def signal_kernel(self, kernel_id, signum):
         self.log.info("Signaled Kernel %s with %s" % (kernel_id, signum))
 
     @kernel_method
-    def restart_kernel(self, kernel_id):
+    def restart_kernel(self, kernel_id, now=False):
         """Restart a kernel by its uuid, keeping the same ports.
 
         Parameters

IPython/kernel/tests/test_kernel.py

@@ -11,104 +11,18 @@
 # Imports
 #-------------------------------------------------------------------------------
 
-import os
-import shutil
 import sys
-import tempfile
-
-from contextlib import contextmanager
-from subprocess import PIPE
 
 import nose.tools as nt
 
-from IPython.kernel import KernelManager
-from IPython.kernel.tests.test_message_spec import execute, flush_channels
 from IPython.testing import decorators as dec, tools as tt
-from IPython.utils import path, py3compat
+from IPython.utils import py3compat
+
+from .utils import new_kernel, kernel, TIMEOUT, assemble_output, execute, flush_channels
 
 #-------------------------------------------------------------------------------
 # Tests
 #-------------------------------------------------------------------------------
-IPYTHONDIR = None
-save_env = None
-save_get_ipython_dir = None
-
-STARTUP_TIMEOUT = 60
-TIMEOUT = 15
-
-def setup():
-    """setup temporary IPYTHONDIR for tests"""
-    global IPYTHONDIR
-    global save_env
-    global save_get_ipython_dir
-
-    IPYTHONDIR = tempfile.mkdtemp()
-
-    save_env = os.environ.copy()
-    os.environ["IPYTHONDIR"] = IPYTHONDIR
-
-    save_get_ipython_dir = path.get_ipython_dir
-    path.get_ipython_dir = lambda : IPYTHONDIR
-
-
-def teardown():
-    path.get_ipython_dir = save_get_ipython_dir
-    os.environ = save_env
-
-    try:
-        shutil.rmtree(IPYTHONDIR)
-    except (OSError, IOError):
-        # no such file
-        pass
-
-
-@contextmanager
-def new_kernel():
-    """start a kernel in a subprocess, and wait for it to be ready
-
-    Returns
-    -------
-    kernel_manager: connected KernelManager instance
-    """
-    KM = KernelManager()
-
-    KM.start_kernel(stdout=PIPE, stderr=PIPE)
-    KC = KM.client()
-    KC.start_channels()
-
-    # wait for kernel to be ready
-    KC.shell_channel.execute("import sys")
-    KC.shell_channel.get_msg(block=True, timeout=STARTUP_TIMEOUT)
-    flush_channels(KC)
-    try:
-        yield KC
-    finally:
-        KC.stop_channels()
-        KM.shutdown_kernel()
-
-
-def assemble_output(iopub):
-    """assemble stdout/err from an execution"""
-    stdout = ''
-    stderr = ''
-    while True:
-        msg = iopub.get_msg(block=True, timeout=1)
-        msg_type = msg['msg_type']
-        content = msg['content']
-        if msg_type == 'status' and content['execution_state'] == 'idle':
-            # idle message signals end of output
-            break
-        elif msg['msg_type'] == 'stream':
-            if content['name'] == 'stdout':
-                stdout = stdout + content['data']
-            elif content['name'] == 'stderr':
-                stderr = stderr + content['data']
-            else:
-                raise KeyError("bad stream: %r" % content['name'])
-        else:
-            # other output, ignored
-            pass
-    return stdout, stderr
 
 
 def _check_mp_mode(kc, expected=False, stream="stdout"):
@@ -123,7 +37,7 @@ def _check_mp_mode(kc, expected=False, stream="stdout"):
 
 def test_simple_print():
     """simple print statement in kernel"""
-    with new_kernel() as kc:
+    with kernel() as kc:
         iopub = kc.iopub_channel
         msg_id, content = execute(kc=kc, code="print ('hi')")
         stdout, stderr = assemble_output(iopub)
@@ -165,7 +79,7 @@ def test_subprocess_print():
 
 def test_subprocess_noprint():
     """mp.Process without print doesn't trigger iostream mp_mode"""
-    with new_kernel() as kc:
+    with kernel() as kc:
         iopub = kc.iopub_channel
 
         np = 5
@@ -210,7 +124,7 @@ def test_subprocess_error():
 
 def test_raw_input():
     """test [raw_]input"""
-    with new_kernel() as kc:
+    with kernel() as kc:
         iopub = kc.iopub_channel
 
         input_f = "input" if py3compat.PY3 else "raw_input"
@@ -232,7 +146,7 @@ def test_raw_input():
 @dec.skipif(py3compat.PY3)
 def test_eval_input():
     """test input() on Python 2"""
-    with new_kernel() as kc:
+    with kernel() as kc:
         iopub = kc.iopub_channel
 
         input_f = "input" if py3compat.PY3 else "raw_input"

IPython/kernel/tests/test_kernelmanager.py

@@ -26,18 +26,11 @@ def _get_ipc_km(self):
     def _run_lifecycle(self, km):
         km.start_kernel(stdout=PIPE, stderr=PIPE)
         self.assertTrue(km.is_alive())
-        km.restart_kernel()
+        km.restart_kernel(now=True)
         self.assertTrue(km.is_alive())
-        # We need a delay here to give the restarting kernel a chance to
-        # restart. Otherwise, the interrupt will kill it, causing the test
-        # suite to hang. The reason it *hangs* is that the shutdown
-        # message for the restart sometimes hasn't been sent to the kernel.
-        # Because linger is oo on the shell channel, the context can't
-        # close until the message is sent to the kernel, which is not dead.
-        time.sleep(1.0)
         km.interrupt_kernel()
         self.assertTrue(isinstance(km, KernelManager))
-        km.shutdown_kernel()
+        km.shutdown_kernel(now=True)
 
     def test_tcp_lifecycle(self):
         km = self._get_tcp_km()

IPython/kernel/tests/test_message_spec.py

@@ -1,7 +1,7 @@
-"""Test suite for our zeromq-based messaging specification.
+"""Test suite for our zeromq-based message specification.
 """
 #-----------------------------------------------------------------------------
-#  Copyright (C) 2010-2011  The IPython Development Team
+#  Copyright (C) 2010  The IPython Development Team
 #
 #  Distributed under the terms of the BSD License.  The full license is in
 #  the file COPYING.txt, distributed as part of this software.
@@ -19,72 +19,21 @@
     HasTraits, TraitError, Bool, Unicode, Dict, Integer, List, Enum, Any,
 )
 
+from .utils import TIMEOUT, start_global_kernel, flush_channels, execute
+
 #-----------------------------------------------------------------------------
-# Global setup and utilities
+# Globals
 #-----------------------------------------------------------------------------
-
-STARTUP_TIMEOUT = 60
-TIMEOUT = 15
+KC = None
 
 def setup():
-    global KM, KC
-    KM = KernelManager()
-    KM.start_kernel(stdout=PIPE, stderr=PIPE)
-    KC = KM.client()
-    KC.start_channels()
-
-    # wait for kernel to be ready
-    try:
-        msg = KC.iopub_channel.get_msg(block=True, timeout=STARTUP_TIMEOUT)
-    except Empty:
-        pass
-    msg_id = KC.kernel_info()
-    KC.get_shell_msg(block=True, timeout=STARTUP_TIMEOUT)
-    flush_channels()
-
-
-def teardown():
-    KC.stop_channels()
-    KM.shutdown_kernel()
-
-
-def flush_channels(kc=None):
-    """flush any messages waiting on the queue"""
-    if kc is None:
-        kc = KC
-    for channel in (kc.shell_channel, kc.iopub_channel):
-        while True:
-            try:
-                msg = channel.get_msg(block=True, timeout=0.1)
-            except Empty:
-                break
-            else:
-                validate_message(msg)
-
-
-def execute(code='', kc=None, **kwargs):
-    """wrapper for doing common steps for validating an execution request"""
-    if kc is None:
-        kc = KC
-    msg_id = kc.execute(code=code, **kwargs)
-    reply = kc.get_shell_msg(timeout=TIMEOUT)
-    validate_message(reply, 'execute_reply', msg_id)
-    busy = kc.get_iopub_msg(timeout=TIMEOUT)
-    validate_message(busy, 'status', msg_id)
-    nt.assert_equal(busy['content']['execution_state'], 'busy')
-
-    if not kwargs.get('silent'):
-        pyin = kc.get_iopub_msg(timeout=TIMEOUT)
-        validate_message(pyin, 'pyin', msg_id)
-        nt.assert_equal(pyin['content']['code'], code)
-
-    return msg_id, reply['content']
+    global KC
+    KC = start_global_kernel()
 
 #-----------------------------------------------------------------------------
-# MSG Spec References
+# Message Spec References
 #-----------------------------------------------------------------------------
 
-
 class Reference(HasTraits):
 
     """

IPython/kernel/tests/test_multikernelmanager.py

@@ -31,20 +31,13 @@ def _run_lifecycle(self, km):
         self.assertTrue(kid in km)
         self.assertTrue(kid in km.list_kernel_ids())
         self.assertEqual(len(km),1)
-        km.restart_kernel(kid)
+        km.restart_kernel(kid, now=True)
         self.assertTrue(km.is_alive(kid))
         self.assertTrue(kid in km.list_kernel_ids())
-        # We need a delay here to give the restarting kernel a chance to
-        # restart. Otherwise, the interrupt will kill it, causing the test
-        # suite to hang. The reason it *hangs* is that the shutdown
-        # message for the restart sometimes hasn't been sent to the kernel.
-        # Because linger is oo on the shell channel, the context can't
-        # close until the message is sent to the kernel, which is not dead.
-        time.sleep(1.0)
         km.interrupt_kernel(kid)
         k = km.get_kernel(kid)
         self.assertTrue(isinstance(k, KernelManager))
-        km.shutdown_kernel(kid)
+        km.shutdown_kernel(kid, now=True)
         self.assertTrue(not kid in km)
 
     def _run_cinfo(self, km, transport, ip):
@@ -63,7 +56,7 @@ def _run_cinfo(self, km, transport, ip):
         self.assertTrue('hb_port' in cinfo)
         stream = km.connect_hb(kid)
         stream.close()
-        km.shutdown_kernel(kid)
+        km.shutdown_kernel(kid, now=True)
 
     def test_tcp_lifecycle(self):
         km = self._get_tcp_km()