Merge pull request #4149 from minrk/fewer-kernels

reuse more kernels in kernel tests

IPython/kernel/tests/utils.py

@@ -0,0 +1,167 b''
+"""utilities for testing IPython kernels"""
+
+#-------------------------------------------------------------------------------
+#  Copyright (C) 2013  The IPython Development Team
+#
+#  Distributed under the terms of the BSD License.  The full license is in
+#  the file COPYING, distributed as part of this software.
+#-------------------------------------------------------------------------------
+
+#-------------------------------------------------------------------------------
+# Imports
+#-------------------------------------------------------------------------------
+
+import atexit
+
+from contextlib import contextmanager
+from subprocess import PIPE
+from Queue import Empty
+
+import nose.tools as nt
+
+from IPython.kernel import KernelManager
+
+#-------------------------------------------------------------------------------
+# Globals
+#-------------------------------------------------------------------------------
+
+STARTUP_TIMEOUT = 60
+TIMEOUT = 15
+
+KM = None
+KC = None
+
+#-------------------------------------------------------------------------------
+# code
+#-------------------------------------------------------------------------------
+
+
+def start_new_kernel():
+    """start a new kernel, and return its Manager and Client"""
+    km = KernelManager()
+    km.start_kernel(stdout=PIPE, stderr=PIPE)
+    kc = km.client()
+    kc.start_channels()
+    
+    msg_id = kc.kernel_info()
+    kc.get_shell_msg(block=True, timeout=STARTUP_TIMEOUT)
+    flush_channels(kc)
+    return km, kc
+
+def flush_channels(kc=None):
+    """flush any messages waiting on the queue"""
+    from .test_message_spec import validate_message
+    
+    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"""
+    from .test_message_spec import validate_message
+    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']
+
+def start_global_kernel():
+    """start the global kernel (if it isn't running) and return its client"""
+    global KM, KC
+    if KM is None:
+        KM, KC = start_new_kernel()
+        atexit.register(stop_global_kernel)
+    return KC
+
+@contextmanager
+def kernel():
+    """Context manager for the global kernel instance
+    
+    Should be used for most kernel tests
+    
+    Returns
+    -------
+    kernel_client: connected KernelClient instance
+    """
+    yield start_global_kernel()
+
+def uses_kernel(test_f):
+    """Decorator for tests that use the global kernel"""
+    def wrapped_test():
+        with kernel() as kc:
+            test_f(kc)
+    wrapped_test.__doc__ = test_f.__doc__
+    wrapped_test.__name__ = test_f.__name__
+    return wrapped_test
+
+def stop_global_kernel():
+    """Stop the global shared kernel instance, if it exists"""
+    global KM, KC
+    KC.stop_channels()
+    KC = None
+    if KM is None:
+        return
+    KM.shutdown_kernel(now=True)
+    KM = None
+
+@contextmanager
+def new_kernel():
+    """Context manager for a new kernel in a subprocess
+    
+    Should only be used for tests where the kernel must not be re-used.
+    
+    Returns
+    -------
+    kernel_client: connected KernelClient instance
+    """
+    km, kc = start_new_kernel()
+    try:
+        yield kc
+    finally:
+        kc.stop_channels()
+        km.shutdown_kernel(now=True)
+
+
+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 += content['data']
+            elif content['name'] == 'stderr':
+                stderr += content['data']
+            else:
+                raise KeyError("bad stream: %r" % content['name'])
+        else:
+            # other output, ignored
+            pass
+    return stdout, stderr
+
+
+

IPython/kernel/multikernelmanager.py

@@ -1,301 +1,301 b''
 """A kernel manager for multiple kernels
 
 Authors:
 
 * Brian Granger
 """
 
 #-----------------------------------------------------------------------------
 #  Copyright (C) 2013  The IPython Development Team
 #
 #  Distributed under the terms of the BSD License.  The full license is in
 #  the file COPYING, distributed as part of this software.
 #-----------------------------------------------------------------------------
 
 #-----------------------------------------------------------------------------
 # Imports
 #-----------------------------------------------------------------------------
 
 from __future__ import absolute_import
 
 import os
 import uuid
 
 import zmq
 
 from IPython.config.configurable import LoggingConfigurable
 from IPython.utils.importstring import import_item
 from IPython.utils.traitlets import (
     Instance, Dict, Unicode, Any, DottedObjectName
 )
 
 #-----------------------------------------------------------------------------
 # Classes
 #-----------------------------------------------------------------------------
 
 class DuplicateKernelError(Exception):
     pass
 
 
 
 def kernel_method(f):
     """decorator for proxying MKM.method(kernel_id) to individual KMs by ID"""
     def wrapped(self, kernel_id, *args, **kwargs):
         # get the kernel
         km = self.get_kernel(kernel_id)
         method = getattr(km, f.__name__)
         # call the kernel's method
         r = method(*args, **kwargs)
         # last thing, call anything defined in the actual class method
         # such as logging messages
         f(self, kernel_id, *args, **kwargs)
         # return the method result
         return r
     return wrapped
 
 
 class MultiKernelManager(LoggingConfigurable):
     """A class for managing multiple kernels."""
     
     kernel_manager_class = DottedObjectName(
         "IPython.kernel.ioloop.IOLoopKernelManager", config=True,
         help="""The kernel manager class.  This is configurable to allow
         subclassing of the KernelManager for customized behavior.
         """
     )
     def _kernel_manager_class_changed(self, name, old, new):
         self.kernel_manager_factory = import_item(new)
     
     kernel_manager_factory = Any(help="this is kernel_manager_class after import")
     def _kernel_manager_factory_default(self):
         return import_item(self.kernel_manager_class)
     
     context = Instance('zmq.Context')
     def _context_default(self):
         return zmq.Context.instance()
 
     connection_dir = Unicode('')
 
     _kernels = Dict()
 
     def list_kernel_ids(self):
         """Return a list of the kernel ids of the active kernels."""
         # Create a copy so we can iterate over kernels in operations
         # that delete keys.
         return list(self._kernels.keys())
 
     def __len__(self):
         """Return the number of running kernels."""
         return len(self.list_kernel_ids())
 
     def __contains__(self, kernel_id):
         return kernel_id in self._kernels
 
     def start_kernel(self, **kwargs):
         """Start a new kernel.
 
         The caller can pick a kernel_id by passing one in as a keyword arg,
         otherwise one will be picked using a uuid.
 
         To silence the kernel's stdout/stderr, call this using::
 
             km.start_kernel(stdout=PIPE, stderr=PIPE)
 
         """
         kernel_id = kwargs.pop('kernel_id', unicode(uuid.uuid4()))
         if kernel_id in self:
             raise DuplicateKernelError('Kernel already exists: %s' % kernel_id)
         # kernel_manager_factory is the constructor for the KernelManager
         # subclass we are using. It can be configured as any Configurable,
         # including things like its transport and ip.
         km = self.kernel_manager_factory(connection_file=os.path.join(
                     self.connection_dir, "kernel-%s.json" % kernel_id),
                     parent=self, autorestart=True, log=self.log
         )
         km.start_kernel(**kwargs)
         self._kernels[kernel_id] = km
         return kernel_id
 
     @kernel_method
     def shutdown_kernel(self, kernel_id, now=False):
         """Shutdown a kernel by its kernel uuid.
 
         Parameters
         ==========
         kernel_id : uuid
             The id of the kernel to shutdown.
         now : bool
             Should the kernel be shutdown forcibly using a signal.
         """
         self.log.info("Kernel shutdown: %s" % kernel_id)
         self.remove_kernel(kernel_id)
 
     def remove_kernel(self, kernel_id):
         """remove a kernel from our mapping.
 
         Mainly so that a kernel can be removed if it is already dead,
         without having to call shutdown_kernel.
 
         The kernel object is returned.
         """
         return self._kernels.pop(kernel_id)
 
     def shutdown_all(self, now=False):
         """Shutdown all kernels."""
         for kid in self.list_kernel_ids():
             self.shutdown_kernel(kid, now=now)
 
     @kernel_method
     def interrupt_kernel(self, kernel_id):
         """Interrupt (SIGINT) the kernel by its uuid.
 
         Parameters
         ==========
         kernel_id : uuid
             The id of the kernel to interrupt.
         """
         self.log.info("Kernel interrupted: %s" % kernel_id)
 
     @kernel_method
     def signal_kernel(self, kernel_id, signum):
         """Sends a signal to the kernel by its uuid.
 
         Note that since only SIGTERM is supported on Windows, this function
         is only useful on Unix systems.
 
         Parameters
         ==========
         kernel_id : uuid
             The id of the kernel to signal.
         """
         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
         ==========
         kernel_id : uuid
             The id of the kernel to interrupt.
         """
         self.log.info("Kernel restarted: %s" % kernel_id)
 
     @kernel_method
     def is_alive(self, kernel_id):
         """Is the kernel alive.
 
         This calls KernelManager.is_alive() which calls Popen.poll on the
         actual kernel subprocess.
 
         Parameters
         ==========
         kernel_id : uuid
             The id of the kernel.
         """
 
     def _check_kernel_id(self, kernel_id):
         """check that a kernel id is valid"""
         if kernel_id not in self:
             raise KeyError("Kernel with id not found: %s" % kernel_id)
 
     def get_kernel(self, kernel_id):
         """Get the single KernelManager object for a kernel by its uuid.
 
         Parameters
         ==========
         kernel_id : uuid
             The id of the kernel.
         """
         self._check_kernel_id(kernel_id)
         return self._kernels[kernel_id]
 
     @kernel_method
     def add_restart_callback(self, kernel_id, callback, event='restart'):
         """add a callback for the KernelRestarter"""
 
     @kernel_method
     def remove_restart_callback(self, kernel_id, callback, event='restart'):
         """remove a callback for the KernelRestarter"""
 
     @kernel_method
     def get_connection_info(self, kernel_id):
         """Return a dictionary of connection data for a kernel.
 
         Parameters
         ==========
         kernel_id : uuid
             The id of the kernel.
 
         Returns
         =======
         connection_dict : dict
             A dict of the information needed to connect to a kernel.
             This includes the ip address and the integer port
             numbers of the different channels (stdin_port, iopub_port,
             shell_port, hb_port).
         """
 
     @kernel_method
     def connect_iopub(self, kernel_id, identity=None):
         """Return a zmq Socket connected to the iopub channel.
 
         Parameters
         ==========
         kernel_id : uuid
             The id of the kernel
         identity : bytes (optional)
             The zmq identity of the socket
 
         Returns
         =======
         stream : zmq Socket or ZMQStream
         """
 
     @kernel_method
     def connect_shell(self, kernel_id, identity=None):
         """Return a zmq Socket connected to the shell channel.
 
         Parameters
         ==========
         kernel_id : uuid
             The id of the kernel
         identity : bytes (optional)
             The zmq identity of the socket
 
         Returns
         =======
         stream : zmq Socket or ZMQStream
         """
 
     @kernel_method
     def connect_stdin(self, kernel_id, identity=None):
         """Return a zmq Socket connected to the stdin channel.
 
         Parameters
         ==========
         kernel_id : uuid
             The id of the kernel
         identity : bytes (optional)
             The zmq identity of the socket
 
         Returns
         =======
         stream : zmq Socket or ZMQStream
         """
 
     @kernel_method
     def connect_hb(self, kernel_id, identity=None):
         """Return a zmq Socket connected to the hb channel.
 
         Parameters
         ==========
         kernel_id : uuid
             The id of the kernel
         identity : bytes (optional)
             The zmq identity of the socket
 
         Returns
         =======
         stream : zmq Socket or ZMQStream
         """

IPython/kernel/tests/test_kernel.py

@@ -1,256 +1,170 b''
 """test the IPython Kernel"""
 
 #-------------------------------------------------------------------------------
 #  Copyright (C) 2013  The IPython Development Team
 #
 #  Distributed under the terms of the BSD License.  The full license is in
 #  the file COPYING, distributed as part of this software.
 #-------------------------------------------------------------------------------
 
 #-------------------------------------------------------------------------------
 # 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"):
     execute(kc=kc, code="import sys")
     flush_channels(kc)
     msg_id, content = execute(kc=kc, code="print (sys.%s._check_mp_mode())" % stream)
     stdout, stderr = assemble_output(kc.iopub_channel)
     nt.assert_equal(eval(stdout.strip()), expected)
 
 
 # printing tests
 
 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)
         nt.assert_equal(stdout, 'hi\n')
         nt.assert_equal(stderr, '')
         _check_mp_mode(kc, expected=False)
 
 
 @dec.knownfailureif(sys.platform == 'win32', "subprocess prints fail on Windows")
 def test_subprocess_print():
     """printing from forked mp.Process"""
     with new_kernel() as kc:
         iopub = kc.iopub_channel
         
         _check_mp_mode(kc, expected=False)
         flush_channels(kc)
         np = 5
         code = '\n'.join([
             "from __future__ import print_function",
             "import multiprocessing as mp",
             "pool = [mp.Process(target=print, args=('hello', i,)) for i in range(%i)]" % np,
             "for p in pool: p.start()",
             "for p in pool: p.join()"
         ])
         
         expected = '\n'.join([
             "hello %s" % i for i in range(np)
         ]) + '\n'
         
         msg_id, content = execute(kc=kc, code=code)
         stdout, stderr = assemble_output(iopub)
         nt.assert_equal(stdout.count("hello"), np, stdout)
         for n in range(np):
             nt.assert_equal(stdout.count(str(n)), 1, stdout)
         nt.assert_equal(stderr, '')
         _check_mp_mode(kc, expected=False)
         _check_mp_mode(kc, expected=False, stream="stderr")
 
 
 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
         code = '\n'.join([
             "import multiprocessing as mp",
             "pool = [mp.Process(target=range, args=(i,)) for i in range(%i)]" % np,
             "for p in pool: p.start()",
             "for p in pool: p.join()"
         ])
         
         msg_id, content = execute(kc=kc, code=code)
         stdout, stderr = assemble_output(iopub)
         nt.assert_equal(stdout, '')
         nt.assert_equal(stderr, '')
 
         _check_mp_mode(kc, expected=False)
         _check_mp_mode(kc, expected=False, stream="stderr")
 
 
 @dec.knownfailureif(sys.platform == 'win32', "subprocess prints fail on Windows")
 def test_subprocess_error():
     """error in mp.Process doesn't crash"""
     with new_kernel() as kc:
         iopub = kc.iopub_channel
         
         code = '\n'.join([
             "import multiprocessing as mp",
             "p = mp.Process(target=int, args=('hi',))",
             "p.start()",
             "p.join()",
         ])
         
         msg_id, content = execute(kc=kc, code=code)
         stdout, stderr = assemble_output(iopub)
         nt.assert_equal(stdout, '')
         nt.assert_true("ValueError" in stderr, stderr)
 
         _check_mp_mode(kc, expected=False)
         _check_mp_mode(kc, expected=False, stream="stderr")
 
 # raw_input tests
 
 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"
         theprompt = "prompt> "
         code = 'print({input_f}("{theprompt}"))'.format(**locals())
         msg_id = kc.execute(code, allow_stdin=True)
         msg = kc.get_stdin_msg(block=True, timeout=TIMEOUT)
         nt.assert_equal(msg['header']['msg_type'], u'input_request')
         content = msg['content']
         nt.assert_equal(content['prompt'], theprompt)
         text = "some text"
         kc.input(text)
         reply = kc.get_shell_msg(block=True, timeout=TIMEOUT)
         nt.assert_equal(reply['content']['status'], 'ok')
         stdout, stderr = assemble_output(iopub)
         nt.assert_equal(stdout, text + "\n")
 
 
 @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"
         theprompt = "prompt> "
         code = 'print(input("{theprompt}"))'.format(**locals())
         msg_id = kc.execute(code, allow_stdin=True)
         msg = kc.get_stdin_msg(block=True, timeout=TIMEOUT)
         nt.assert_equal(msg['header']['msg_type'], u'input_request')
         content = msg['content']
         nt.assert_equal(content['prompt'], theprompt)
         kc.input("1+1")
         reply = kc.get_shell_msg(block=True, timeout=TIMEOUT)
         nt.assert_equal(reply['content']['status'], 'ok')
         stdout, stderr = assemble_output(iopub)
         nt.assert_equal(stdout, "2\n")
 
 
 def test_help_output():
     """ipython kernel --help-all works"""
     tt.help_all_output_test('kernel')
 

IPython/kernel/tests/test_kernelmanager.py

@@ -1,61 +1,54 b''
 """Tests for the notebook kernel and session manager"""
 
 from subprocess import PIPE
 import time
 from unittest import TestCase
 
 from IPython.testing import decorators as dec
 
 from IPython.config.loader import Config
 from IPython.kernel import KernelManager
 
 class TestKernelManager(TestCase):
 
     def _get_tcp_km(self):
         c = Config()
         km = KernelManager(config=c)
         return km
 
     def _get_ipc_km(self):
         c = Config()
         c.KernelManager.transport = 'ipc'
         c.KernelManager.ip = 'test'
         km = KernelManager(config=c)
         return km
 
     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()
         self._run_lifecycle(km)
 
     @dec.skip_win32
     def test_ipc_lifecycle(self):
         km = self._get_ipc_km()
         self._run_lifecycle(km)
     
     def test_get_connect_info(self):
         km = self._get_tcp_km()
         cinfo = km.get_connection_info()
         keys = sorted(cinfo.keys())
         expected = sorted([
             'ip', 'transport',
             'hb_port', 'shell_port', 'stdin_port', 'iopub_port', 'control_port',
             'key', 'signature_scheme',
         ])
         self.assertEqual(keys, expected)
 

IPython/kernel/tests/test_message_spec.py

@@ -1,484 +1,433 b''
-"""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.
 #-----------------------------------------------------------------------------
 
 import re
 from subprocess import PIPE
 from Queue import Empty
 
 import nose.tools as nt
 
 from IPython.kernel import KernelManager
 
 from IPython.utils.traitlets import (
     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):
 
     """
     Base class for message spec specification testing.
 
     This class is the core of the message specification test.  The
     idea is that child classes implement trait attributes for each
     message keys, so that message keys can be tested against these
     traits using :meth:`check` method.
 
     """
 
     def check(self, d):
         """validate a dict against our traits"""
         for key in self.trait_names():
             nt.assert_in(key, d)
             # FIXME: always allow None, probably not a good idea
             if d[key] is None:
                 continue
             try:
                 setattr(self, key, d[key])
             except TraitError as e:
                 nt.assert_true(False, str(e))
 
 
 class RMessage(Reference):
     msg_id = Unicode()
     msg_type = Unicode()
     header = Dict()
     parent_header = Dict()
     content = Dict()
 
 class RHeader(Reference):
     msg_id = Unicode()
     msg_type = Unicode()
     session = Unicode()
     username = Unicode()
 
 class RContent(Reference):
     status = Enum((u'ok', u'error'))
 
 
 class ExecuteReply(Reference):
     execution_count = Integer()
     status = Enum((u'ok', u'error'))
     
     def check(self, d):
         Reference.check(self, d)
         if d['status'] == 'ok':
             ExecuteReplyOkay().check(d)
         elif d['status'] == 'error':
             ExecuteReplyError().check(d)
 
 
 class ExecuteReplyOkay(Reference):
     payload = List(Dict)
     user_variables = Dict()
     user_expressions = Dict()
 
 
 class ExecuteReplyError(Reference):
     ename = Unicode()
     evalue = Unicode()
     traceback = List(Unicode)
 
 
 class OInfoReply(Reference):
     name = Unicode()
     found = Bool()
     ismagic = Bool()
     isalias = Bool()
     namespace = Enum((u'builtin', u'magics', u'alias', u'Interactive'))
     type_name = Unicode()
     string_form = Unicode()
     base_class = Unicode()
     length = Integer()
     file = Unicode()
     definition = Unicode()
     argspec = Dict()
     init_definition = Unicode()
     docstring = Unicode()
     init_docstring = Unicode()
     class_docstring = Unicode()
     call_def = Unicode()
     call_docstring = Unicode()
     source = Unicode()
     
     def check(self, d):
         Reference.check(self, d)
         if d['argspec'] is not None:
             ArgSpec().check(d['argspec'])
 
 
 class ArgSpec(Reference):
     args = List(Unicode)
     varargs = Unicode()
     varkw = Unicode()
     defaults = List()
 
 
 class Status(Reference):
     execution_state = Enum((u'busy', u'idle', u'starting'))
 
 
 class CompleteReply(Reference):
     matches = List(Unicode)
 
 
 def Version(num, trait=Integer):
     return List(trait, default_value=[0] * num, minlen=num, maxlen=num)
 
 
 class KernelInfoReply(Reference):
 
     protocol_version = Version(2)
     ipython_version = Version(4, Any)
     language_version = Version(3)
     language = Unicode()
 
     def _ipython_version_changed(self, name, old, new):
         for v in new:
             assert isinstance(v, int) or isinstance(v, basestring), \
             'expected int or string as version component, got {0!r}'.format(v)
 
 
 # IOPub messages
 
 class PyIn(Reference):
     code = Unicode()
     execution_count = Integer()
 
 
 PyErr = ExecuteReplyError
 
 
 class Stream(Reference):
     name = Enum((u'stdout', u'stderr'))
     data = Unicode()
 
 
 mime_pat = re.compile(r'\w+/\w+')
 
 class DisplayData(Reference):
     source = Unicode()
     metadata = Dict()
     data = Dict()
     def _data_changed(self, name, old, new):
         for k,v in new.iteritems():
             assert mime_pat.match(k)
             nt.assert_is_instance(v, basestring)
 
 
 class PyOut(Reference):
     execution_count = Integer()
     data = Dict()
     def _data_changed(self, name, old, new):
         for k,v in new.iteritems():
             assert mime_pat.match(k)
             nt.assert_is_instance(v, basestring)
 
 
 references = {
     'execute_reply' : ExecuteReply(),
     'object_info_reply' : OInfoReply(),
     'status' : Status(),
     'complete_reply' : CompleteReply(),
     'kernel_info_reply': KernelInfoReply(),
     'pyin' : PyIn(),
     'pyout' : PyOut(),
     'pyerr' : PyErr(),
     'stream' : Stream(),
     'display_data' : DisplayData(),
 }
 """
 Specifications of `content` part of the reply messages.
 """
 
 
 def validate_message(msg, msg_type=None, parent=None):
     """validate a message
     
     This is a generator, and must be iterated through to actually
     trigger each test.
     
     If msg_type and/or parent are given, the msg_type and/or parent msg_id
     are compared with the given values.
     """
     RMessage().check(msg)
     if msg_type:
         nt.assert_equal(msg['msg_type'], msg_type)
     if parent:
         nt.assert_equal(msg['parent_header']['msg_id'], parent)
     content = msg['content']
     ref = references[msg['msg_type']]
     ref.check(content)
 
 
 #-----------------------------------------------------------------------------
 # Tests
 #-----------------------------------------------------------------------------
 
 # Shell channel
 
 def test_execute():
     flush_channels()
     
     msg_id = KC.execute(code='x=1')
     reply = KC.get_shell_msg(timeout=TIMEOUT)
     validate_message(reply, 'execute_reply', msg_id)
 
 
 def test_execute_silent():
     flush_channels()
     msg_id, reply = execute(code='x=1', silent=True)
     
     # flush status=idle
     status = KC.iopub_channel.get_msg(timeout=TIMEOUT)
     validate_message(status, 'status', msg_id)
     nt.assert_equal(status['content']['execution_state'], 'idle')
 
     nt.assert_raises(Empty, KC.iopub_channel.get_msg, timeout=0.1)
     count = reply['execution_count']
     
     msg_id, reply = execute(code='x=2', silent=True)
     
     # flush status=idle
     status = KC.iopub_channel.get_msg(timeout=TIMEOUT)
     validate_message(status, 'status', msg_id)
     nt.assert_equal(status['content']['execution_state'], 'idle')
     
     nt.assert_raises(Empty, KC.iopub_channel.get_msg, timeout=0.1)
     count_2 = reply['execution_count']
     nt.assert_equal(count_2, count)
 
 
 def test_execute_error():
     flush_channels()
     
     msg_id, reply = execute(code='1/0')
     nt.assert_equal(reply['status'], 'error')
     nt.assert_equal(reply['ename'], 'ZeroDivisionError')
     
     pyerr = KC.iopub_channel.get_msg(timeout=TIMEOUT)
     validate_message(pyerr, 'pyerr', msg_id)
 
 
 def test_execute_inc():
     """execute request should increment execution_count"""
     flush_channels()
 
     msg_id, reply = execute(code='x=1')
     count = reply['execution_count']
     
     flush_channels()
     
     msg_id, reply = execute(code='x=2')
     count_2 = reply['execution_count']
     nt.assert_equal(count_2, count+1)
 
 
 def test_user_variables():
     flush_channels()
 
     msg_id, reply = execute(code='x=1', user_variables=['x'])
     user_variables = reply['user_variables']
     nt.assert_equal(user_variables, {u'x': {
         u'status': u'ok',
         u'data': {u'text/plain': u'1'},
         u'metadata': {},
     }})
 
 
 def test_user_variables_fail():
     flush_channels()
 
     msg_id, reply = execute(code='x=1', user_variables=['nosuchname'])
     user_variables = reply['user_variables']
     foo = user_variables['nosuchname']
     nt.assert_equal(foo['status'], 'error')
     nt.assert_equal(foo['ename'], 'KeyError')
 
 
 def test_user_expressions():
     flush_channels()
 
     msg_id, reply = execute(code='x=1', user_expressions=dict(foo='x+1'))
     user_expressions = reply['user_expressions']
     nt.assert_equal(user_expressions, {u'foo': {
         u'status': u'ok',
         u'data': {u'text/plain': u'2'},
         u'metadata': {},
     }})
 
 
 def test_user_expressions_fail():
     flush_channels()
 
     msg_id, reply = execute(code='x=0', user_expressions=dict(foo='nosuchname'))
     user_expressions = reply['user_expressions']
     foo = user_expressions['foo']
     nt.assert_equal(foo['status'], 'error')
     nt.assert_equal(foo['ename'], 'NameError')
 
 
 def test_oinfo():
     flush_channels()
 
     msg_id = KC.object_info('a')
     reply = KC.get_shell_msg(timeout=TIMEOUT)
     validate_message(reply, 'object_info_reply', msg_id)
 
 
 def test_oinfo_found():
     flush_channels()
 
     msg_id, reply = execute(code='a=5')
     
     msg_id = KC.object_info('a')
     reply = KC.get_shell_msg(timeout=TIMEOUT)
     validate_message(reply, 'object_info_reply', msg_id)
     content = reply['content']
     assert content['found']
     argspec = content['argspec']
     nt.assert_is(argspec, None)
 
 
 def test_oinfo_detail():
     flush_channels()
 
     msg_id, reply = execute(code='ip=get_ipython()')
     
     msg_id = KC.object_info('ip.object_inspect', detail_level=2)
     reply = KC.get_shell_msg(timeout=TIMEOUT)
     validate_message(reply, 'object_info_reply', msg_id)
     content = reply['content']
     assert content['found']
     argspec = content['argspec']
     nt.assert_is_instance(argspec, dict, "expected non-empty argspec dict, got %r" % argspec)
     nt.assert_equal(argspec['defaults'], [0])
 
 
 def test_oinfo_not_found():
     flush_channels()
 
     msg_id = KC.object_info('dne')
     reply = KC.get_shell_msg(timeout=TIMEOUT)
     validate_message(reply, 'object_info_reply', msg_id)
     content = reply['content']
     nt.assert_false(content['found'])
 
 
 def test_complete():
     flush_channels()
 
     msg_id, reply = execute(code="alpha = albert = 5")
     
     msg_id = KC.complete('al', 'al', 2)
     reply = KC.get_shell_msg(timeout=TIMEOUT)
     validate_message(reply, 'complete_reply', msg_id)
     matches = reply['content']['matches']
     for name in ('alpha', 'albert'):
         nt.assert_in(name, matches)
 
 
 def test_kernel_info_request():
     flush_channels()
 
     msg_id = KC.kernel_info()
     reply = KC.get_shell_msg(timeout=TIMEOUT)
     validate_message(reply, 'kernel_info_reply', msg_id)
 
 
 # IOPub channel
 
 
 def test_stream():
     flush_channels()
 
     msg_id, reply = execute("print('hi')")
 
     stdout = KC.iopub_channel.get_msg(timeout=TIMEOUT)
     validate_message(stdout, 'stream', msg_id)
     content = stdout['content']
     nt.assert_equal(content['name'], u'stdout')
     nt.assert_equal(content['data'], u'hi\n')
 
 
 def test_display_data():
     flush_channels()
 
     msg_id, reply = execute("from IPython.core.display import display; display(1)")
     
     display = KC.iopub_channel.get_msg(timeout=TIMEOUT)
     validate_message(display, 'display_data', parent=msg_id)
     data = display['content']['data']
     nt.assert_equal(data['text/plain'], u'1')
 

IPython/kernel/tests/test_multikernelmanager.py

@@ -1,85 +1,78 b''
 """Tests for the notebook kernel and session manager."""
 
 from subprocess import PIPE
 import time
 from unittest import TestCase
 
 from IPython.testing import decorators as dec
 
 from IPython.config.loader import Config
 from IPython.utils.localinterfaces import LOCALHOST
 from IPython.kernel import KernelManager
 from IPython.kernel.multikernelmanager import MultiKernelManager
 
 class TestKernelManager(TestCase):
 
     def _get_tcp_km(self):
         c = Config()
         km = MultiKernelManager(config=c)
         return km
 
     def _get_ipc_km(self):
         c = Config()
         c.KernelManager.transport = 'ipc'
         c.KernelManager.ip = 'test'
         km = MultiKernelManager(config=c)
         return km
 
     def _run_lifecycle(self, km):
         kid = km.start_kernel(stdout=PIPE, stderr=PIPE)
         self.assertTrue(km.is_alive(kid))
         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):
         kid = km.start_kernel(stdout=PIPE, stderr=PIPE)
         k = km.get_kernel(kid)
         cinfo = km.get_connection_info(kid)
         self.assertEqual(transport, cinfo['transport'])
         self.assertEqual(ip, cinfo['ip'])
         self.assertTrue('stdin_port' in cinfo)
         self.assertTrue('iopub_port' in cinfo)
         stream = km.connect_iopub(kid)
         stream.close()
         self.assertTrue('shell_port' in cinfo)
         stream = km.connect_shell(kid)
         stream.close()
         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()
         self._run_lifecycle(km)
     
     def test_tcp_cinfo(self):
         km = self._get_tcp_km()
         self._run_cinfo(km, 'tcp', LOCALHOST)
 
     @dec.skip_win32
     def test_ipc_lifecycle(self):
         km = self._get_ipc_km()
         self._run_lifecycle(km)
     
     @dec.skip_win32
     def test_ipc_cinfo(self):
         km = self._get_ipc_km()
         self._run_cinfo(km, 'ipc', 'test')