IPython/zmq/parallel/tests/clienttest.py

import sys
import tempfile
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 import error
from IPython.zmq.parallel.client import Client
from IPython.zmq.parallel.ipcluster import launch_process
from IPython.zmq.parallel.entry_point import select_random_ports
from IPython.zmq.parallel.tests import processes,add_engine

# simple tasks for use in apply tests

def segfault():
    """this will 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 len(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 assertRaisesRemote(self, etype, f, *args, **kwargs):
        try:
            try:
                f(*args, **kwargs)
            except error.CompositeError as e:
                e.raise_exception()
        except error.RemoteError as e:
            self.assertEquals(etype.__name__, e.ename, "Should have raised %r, but raised %r"%(e.ename, etype.__name__))
        else:
            self.fail("should have raised a RemoteError")
            
    def setUp(self):
        BaseZMQTestCase.setUp(self)
        self.client = self.connect_client()
        self.base_engine_count=len(self.client.ids)
        self.engines=[]
    
    def tearDown(self):
        
        # close fds:
        for e in filter(lambda e: e.poll() is not None, processes):
            processes.remove(e)
        
        self.client.close()
        BaseZMQTestCase.tearDown(self)
        # this will be superfluous when pyzmq merges PR #88
        self.context.term()
        print tempfile.TemporaryFile().fileno(),
        sys.stdout.flush()