IPython/zmq/parallel/tests/clienttest.py

import time
from signal import SIGINT
from multiprocessing import Process

from zmq.tests import BaseZMQTestCase

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 cluster_logs,add_engine


class ClusterTestCase(BaseZMQTestCase):
    
    def add_engines(self, n=1):
        """add multiple engines to our cluster"""
        for i in range(n):
            self.engines.append(add_engine())
    
    def wait_on_engines(self):
        """wait for our engines to connect."""
        while len(self.client.ids) < len(self.engines)+self.base_engine_count:
            time.sleep(0.1)
    
    def start_cluster(self, n=1):
        """start a cluster"""
        raise NotImplementedError("Don't use this anymore")
        rport = select_random_ports(1)[0]
        args = [ '--regport', str(rport), '--ip', '127.0.0.1' ]
        cp = launch_process('controller', args)
        eps = [ launch_process('engine', args+['--ident', 'engine-%i'%i]) for i in range(n) ]
        return rport, args, cp, eps
    
    def connect_client(self, port=None):
        """connect a client with my Context, and track its sockets for cleanup"""
        if port is None:
            port = cluster_logs['regport']
            c = Client('tcp://127.0.0.1:%i'%port,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)