IPython/zmq/parallel/tests/clienttest.py

import time
from signal import SIGINT
from multiprocessing import Process

from nose import SkipTest

from zmq.tests import BaseZMQTestCase

from IPython.external.decorator import decorator

from IPython.zmq.parallel.ipcluster import launch_process
from IPython.zmq.parallel.entry_point import select_random_ports
from IPython.zmq.parallel.client import Client
from IPython.zmq.parallel.tests import processes,add_engine

# simple tasks for use in apply tests

def segfault():
    """"""
    import ctypes
    ctypes.memset(-1,0,1)

def wait(n):
    """sleep for a time"""
    import time
    time.sleep(n)
    return n

def raiser(eclass):
    """raise an exception"""
    raise eclass()

# test decorator for skipping tests when libraries are unavailable
def skip_without(*names):
    """skip a test if some names are not importable"""
    @decorator
    def skip_without_names(f, *args, **kwargs):
        """decorator to skip tests in the absence of numpy."""
        for name in names:
            try:
                __import__(name)
            except ImportError:
                raise SkipTest
        return f(*args, **kwargs)
    return skip_without_names


class ClusterTestCase(BaseZMQTestCase):
    
    def add_engines(self, n=1, block=True):
        """add multiple engines to our cluster"""
        for i in range(n):
            self.engines.append(add_engine())
        if block:
            self.wait_on_engines()
    
    def wait_on_engines(self, timeout=5):
        """wait for our engines to connect."""
        n = len(self.engines)+self.base_engine_count
        tic = time.time()
        while time.time()-tic < timeout and len(self.client.ids) < n:
            time.sleep(0.1)
        
        assert not self.client.ids < n, "waiting for engines timed out"
    
    def connect_client(self):
        """connect a client with my Context, and track its sockets for cleanup"""
        c = Client(profile='iptest',context=self.context)
        for name in filter(lambda n:n.endswith('socket'), dir(c)):
            self.sockets.append(getattr(c, name))
        return c
    
    def setUp(self):
        BaseZMQTestCase.setUp(self)
        self.client = self.connect_client()
        self.base_engine_count=len(self.client.ids)
        self.engines=[]
    
    def tearDown(self):
        [ e.terminate() for e in filter(lambda e: e.poll() is None, self.engines) ]
        # while len(self.client.ids) > self.base_engine_count:
        #     time.sleep(.1)
        del self.engines
        BaseZMQTestCase.tearDown(self)