<s> import os <EOL> import re <EOL> import pty <EOL> import fcntl <EOL> import contextlib <EOL> import subprocess <EOL> import signal <EOL> import math <EOL> from time import time <EOL> from . import basecase <EOL> DEFAULT_CQLSH_PROMPT = '<STR_LIT>' <EOL> DEFAULT_CQLSH_TERM = '<STR_LIT>' <EOL> cqlshlog = basecase . cqlshlog <EOL> def set_controlling_pty ( master , slave ) : <EOL> os . setsid ( ) <EOL> os . close ( master ) <EOL> for i in range ( <NUM_LIT:3> ) : <EOL> os . dup2 ( slave , i ) <EOL> if slave > <NUM_LIT:2> : <EOL> os . close ( slave ) <EOL> os . close ( os . open ( os . ttyname ( <NUM_LIT:1> ) , os . O_RDWR ) ) <EOL> def set_nonblocking ( fd ) : <EOL> flags = fcntl . fcntl ( fd , fcntl . F_GETFL ) <EOL> fcntl . fcntl ( fd , fcntl . F_SETFL , flags | os . O_NONBLOCK ) <EOL> @ contextlib . contextmanager <EOL> def raising_signal ( signum , exc ) : <EOL> """<STR_LIT>""" <EOL> def raiser ( signum , frames ) : <EOL> raise exc ( ) <EOL> oldhandlr = signal . signal ( signum , raiser ) <EOL> try : <EOL> yield <EOL> finally : <EOL> signal . signal ( signum , oldhandlr ) <EOL> class TimeoutError ( Exception ) : <EOL> pass <EOL> @ contextlib . contextmanager <EOL> def timing_out_itimer ( seconds ) : <EOL> if seconds is None : <EOL> yield <EOL> return <EOL> with raising_signal ( signal . SIGALRM , TimeoutError ) : <EOL> oldval , oldint = signal . getitimer ( signal . ITIMER_REAL ) <EOL> if oldval != <NUM_LIT:0.0> : <EOL> raise RuntimeError ( "<STR_LIT>" ) <EOL> signal . setitimer ( signal . ITIMER_REAL , seconds ) <EOL> try : <EOL> yield <EOL> finally : <EOL> signal . setitimer ( signal . ITIMER_REAL , <NUM_LIT:0> ) <EOL> @ contextlib . contextmanager <EOL> def timing_out_alarm ( seconds ) : <EOL> if seconds is None : <EOL> yield <EOL> return <EOL> with raising_signal ( signal . SIGALRM , TimeoutError ) : <EOL> oldval = signal . alarm ( int ( math . ceil ( seconds ) ) ) <EOL> if oldval != <NUM_LIT:0> : <EOL> signal . alarm ( oldval ) <EOL> raise RuntimeError ( "<STR_LIT>" ) <EOL> try : <EOL> yield <EOL> finally : <EOL> signal . alarm ( <NUM_LIT:0> ) <EOL> if hasattr ( signal , '<STR_LIT>' ) : <EOL> timing_out = timing_out_itimer <EOL> else : <EOL> timing_out = timing_out_alarm <EOL> def noop ( * a ) : <EOL> pass <EOL> class ProcRunner : <EOL> def __init__ ( self , path , tty = True , env = None , args = ( ) ) : <EOL> self . exe_path = path <EOL> self . args = args <EOL> self . tty = bool ( tty ) <EOL> if env is None : <EOL> env = { } <EOL> self . env = env <EOL> self . readbuf = '<STR_LIT>' <EOL> self . start_proc ( ) <EOL> def start_proc ( self ) : <EOL> preexec = noop <EOL> stdin = stdout = stderr = None <EOL> if self . tty : <EOL> masterfd , slavefd = pty . openpty ( ) <EOL> preexec = lambda : set_controlling_pty ( masterfd , slavefd ) <EOL> else : <EOL> stdin = stdout = subprocess . PIPE <EOL> stderr = subprocess . STDOUT <EOL> cqlshlog . info ( "<STR_LIT>" <EOL> % ( self . exe_path , self . args , self . env ) ) <EOL> self . proc = subprocess . Popen ( ( self . exe_path , ) + tuple ( self . args ) , <EOL> env = self . env , preexec_fn = preexec , <EOL> stdin = stdin , stdout = stdout , stderr = stderr , <EOL> close_fds = False ) <EOL> if self . tty : <EOL> os . close ( slavefd ) <EOL> self . childpty = masterfd <EOL> self . send = self . send_tty <EOL> self . read = self . read_tty <EOL> else : <EOL> self . send = self . send_pipe <EOL> self . read = self . read_pipe <EOL> def close ( self ) : <EOL> cqlshlog . info ( "<STR_LIT>" % ( self . exe_path , ) ) <EOL> if self . tty : <EOL> os . close ( self . childpty ) <EOL> else : <EOL> self . proc . stdin . close ( ) <EOL> cqlshlog . debug ( "<STR_LIT>" ) <EOL> return self . proc . wait ( ) <EOL> def send_tty ( self , data ) : <EOL> os . write ( self . childpty , data ) <EOL> def send_pipe ( self , data ) : <EOL> self . proc . stdin . write ( data ) <EOL> def read_tty ( self , blksize ) : <EOL> return os . read ( self . childpty , blksize ) <EOL> def read_pipe ( self , blksize ) : <EOL> return self . proc . stdout . read ( blksize ) <EOL> def read_until ( self , until , blksize = <NUM_LIT> , timeout = None , flags = <NUM_LIT:0> ) : <EOL> if not isinstance ( until , re . _pattern_type ) : <EOL> until = re . compile ( until , flags ) <EOL> got = self . readbuf <EOL> self . readbuf = '<STR_LIT>' <EOL> with timing_out ( timeout ) : <EOL> while True : <EOL> val = self . read ( blksize ) <EOL> cqlshlog . debug ( "<STR_LIT>" % ( val , ) ) <EOL> if val == '<STR_LIT>' : <EOL> raise EOFError ( "<STR_LIT>" % ( until . pattern , ) ) <EOL> got += val <EOL> m = until . search ( got ) <EOL> if m is not None : <EOL> self . readbuf = got [ m . end ( ) : ] <EOL> got = got [ : m . end ( ) ] <EOL> return got <EOL> def read_lines ( self , numlines , blksize = <NUM_LIT> , timeout = None ) : <EOL> lines = [ ] <EOL> with timing_out ( timeout ) : <EOL> for n in range ( numlines ) : <EOL> lines . append ( self . read_until ( '<STR_LIT:\n>' , blksize = blksize ) ) <EOL> return lines <EOL> def read_up_to_timeout ( self , timeout , blksize = <NUM_LIT> ) : <EOL> got = self . readbuf <EOL> self . readbuf = '<STR_LIT>' <EOL> curtime = time ( ) <EOL> stoptime = curtime + timeout <EOL> while curtime < stoptime : <EOL> try : <EOL> with timing_out ( stoptime - curtime ) : <EOL> stuff = self . read ( blksize ) <EOL> except TimeoutError : <EOL> break <EOL> cqlshlog . debug ( "<STR_LIT>" % ( stuff , ) ) <EOL> if stuff == '<STR_LIT>' : <EOL> break <EOL> got += stuff <EOL> curtime = time ( ) <EOL> return got <EOL> class CqlshRunner ( ProcRunner ) : <EOL> def __init__ ( self , path = None , host = None , port = None , keyspace = None , cqlver = None , <EOL> args = ( ) , prompt = DEFAULT_CQLSH_PROMPT , env = None , ** kwargs ) : <EOL> if path is None : <EOL> path = basecase . path_to_cqlsh <EOL> if host is None : <EOL> host = basecase . TEST_HOST <EOL> if port is None : <EOL> port = basecase . TEST_PORT <EOL> if env is None : <EOL> env = { } <EOL> env . setdefault ( '<STR_LIT>' , '<STR_LIT>' ) <EOL> env . setdefault ( '<STR_LIT>' , os . environ . get ( '<STR_LIT>' , '<STR_LIT>' ) ) <EOL> env . setdefault ( '<STR_LIT>' , os . environ . get ( '<STR_LIT>' , '<STR_LIT>' ) ) <EOL> args = tuple ( args ) + ( host , str ( port ) ) <EOL> if cqlver is not None : <EOL> args += ( '<STR_LIT>' , str ( cqlver ) ) <EOL> if keyspace is not None : <EOL> args += ( '<STR_LIT>' , keyspace ) <EOL> self . keyspace = keyspace <EOL> ProcRunner . __init__ ( self , path , args = args , env = env , ** kwargs ) <EOL> self . prompt = prompt <EOL> if self . prompt is None : <EOL> self . output_header = '<STR_LIT>' <EOL> else : <EOL> self . output_header = self . read_to_next_prompt ( ) <EOL> def read_to_next_prompt ( self ) : <EOL> return self . read_until ( self . prompt , timeout = <NUM_LIT> ) <EOL> def read_up_to_timeout ( self , timeout , blksize = <NUM_LIT> ) : <EOL> output = ProcRunner . read_up_to_timeout ( self , timeout , blksize = blksize ) <EOL> output = output . replace ( '<STR_LIT>' , '<STR_LIT>' ) <EOL> output = output . replace ( '<STR_LIT:\r>' , '<STR_LIT>' ) <EOL> return output <EOL> def cmd_and_response ( self , cmd ) : <EOL> self . send ( cmd + '<STR_LIT:\n>' ) <EOL> output = self . read_to_next_prompt ( ) <EOL> output = output . replace ( '<STR_LIT>' , '<STR_LIT>' ) <EOL> output = output . replace ( '<STR_LIT:\r>' , '<STR_LIT>' ) <EOL> output = output . replace ( '<STR_LIT>' , '<STR_LIT>' ) <EOL> if self . tty : <EOL> echo , output = output . split ( '<STR_LIT:\n>' , <NUM_LIT:1> ) <EOL> assert echo == cmd , "<STR_LIT>" % ( echo , cmd ) <EOL> try : <EOL> output , promptline = output . rsplit ( '<STR_LIT:\n>' , <NUM_LIT:1> ) <EOL> except ValueError : <EOL> promptline = output <EOL> output = '<STR_LIT>' <EOL> assert re . match ( self . prompt , '<STR_LIT:\n>' + promptline ) , '<STR_LIT>' % ( promptline , self . prompt ) <EOL> return output + '<STR_LIT:\n>' <EOL> def run_cqlsh ( ** kwargs ) : <EOL> return contextlib . closing ( CqlshRunner ( ** kwargs ) ) <EOL> def call_cqlsh ( ** kwargs ) : <EOL> kwargs . setdefault ( '<STR_LIT>' , None ) <EOL> proginput = kwargs . pop ( '<STR_LIT:input>' , '<STR_LIT>' ) <EOL> kwargs [ '<STR_LIT>' ] = False <EOL> c = CqlshRunner ( ** kwargs ) <EOL> output , _ = c . proc . communicate ( proginput ) <EOL> result = c . close ( ) <EOL> return output , result </s>
