<s> import os <EOL> import time <EOL> import uuid <EOL> import shutil <EOL> import socket <EOL> import tempfile <EOL> from contextlib import contextmanager <EOL> from subprocess import Popen , PIPE <EOL> from rb . cluster import Cluster <EOL> devnull = open ( os . devnull , '<STR_LIT>' ) <EOL> class Server ( object ) : <EOL> def __init__ ( self , cl , socket_path ) : <EOL> self . _cl = cl <EOL> self . socket_path = socket_path <EOL> def test_connection ( self ) : <EOL> try : <EOL> s = socket . socket ( socket . AF_UNIX , socket . SOCK_STREAM ) <EOL> s . connect ( self . socket_path ) <EOL> except IOError : <EOL> return False <EOL> return True <EOL> def signal_stop ( self ) : <EOL> if self . _cl is not None : <EOL> self . _cl . kill ( ) <EOL> def close ( self ) : <EOL> if self . _cl is not None : <EOL> self . signal_stop ( ) <EOL> self . _cl . wait ( ) <EOL> self . _cl = None <EOL> try : <EOL> os . remove ( self . socket_path ) <EOL> except OSError : <EOL> pass <EOL> class TestSetup ( object ) : <EOL> """<STR_LIT>""" <EOL> def __init__ ( self , servers = <NUM_LIT:4> , databases_each = <NUM_LIT:8> , <EOL> server_executable = '<STR_LIT>' ) : <EOL> self . _fd_dir = tempfile . mkdtemp ( ) <EOL> self . databases_each = databases_each <EOL> self . server_executable = server_executable <EOL> self . servers = [ ] <EOL> for server in xrange ( servers ) : <EOL> self . spawn_server ( ) <EOL> def __enter__ ( self ) : <EOL> return self <EOL> def __exit__ ( self , exc_type , exc_value , tb ) : <EOL> self . close ( ) <EOL> def make_cluster ( self ) : <EOL> """<STR_LIT>""" <EOL> self . wait_for_servers ( ) <EOL> hosts = [ ] <EOL> host_id = <NUM_LIT:0> <EOL> for server in self . servers : <EOL> for x in xrange ( self . databases_each ) : <EOL> hosts . append ( { <EOL> '<STR_LIT>' : host_id , <EOL> '<STR_LIT>' : server . socket_path , <EOL> '<STR_LIT>' : x , <EOL> } ) <EOL> host_id += <NUM_LIT:1> <EOL> return Cluster ( hosts ) <EOL> def spawn_server ( self ) : <EOL> """<STR_LIT>""" <EOL> socket_path = os . path . join ( self . _fd_dir , str ( uuid . uuid4 ( ) ) ) <EOL> cl = Popen ( [ self . server_executable , '<STR_LIT:->' ] , stdin = PIPE , <EOL> stdout = devnull ) <EOL> cl . stdin . write ( '''<STR_LIT>''' % { <EOL> '<STR_LIT:path>' : socket_path , <EOL> '<STR_LIT>' : self . databases_each , <EOL> } ) <EOL> cl . stdin . flush ( ) <EOL> cl . stdin . close ( ) <EOL> self . servers . append ( Server ( cl , socket_path ) ) <EOL> def wait_for_servers ( self , timeout = <NUM_LIT:10> ) : <EOL> """<STR_LIT>""" <EOL> unconnected_servers = dict ( ( x . socket_path , x ) <EOL> for x in self . servers ) <EOL> now = time . time ( ) <EOL> while unconnected_servers : <EOL> for server in unconnected_servers . itervalues ( ) : <EOL> if server . test_connection ( ) : <EOL> unconnected_servers . pop ( server . socket_path , None ) <EOL> break <EOL> if time . time ( ) > now + timeout : <EOL> return False <EOL> if unconnected_servers : <EOL> time . sleep ( <NUM_LIT> ) <EOL> return True <EOL> def close ( self ) : <EOL> """<STR_LIT>""" <EOL> for server in self . servers : <EOL> server . signal_stop ( ) <EOL> for server in self . servers : <EOL> server . close ( ) <EOL> try : <EOL> shutil . rmtree ( self . _fd_dir ) <EOL> except ( OSError , IOError ) : <EOL> pass <EOL> def __del__ ( self ) : <EOL> try : <EOL> self . close ( ) <EOL> except Exception : <EOL> pass <EOL> @ contextmanager <EOL> def make_test_cluster ( * args , ** kwargs ) : <EOL> """<STR_LIT>""" <EOL> with TestSetup ( * args , ** kwargs ) as ts : <EOL> cluster = ts . make_cluster ( ) <EOL> try : <EOL> yield cluster <EOL> finally : <EOL> cluster . disconnect_pools ( ) </s>
