def start(self):
'''Launch and wait for the process to initialize.'''
print 'starting', self.cmd
self.proc = Popen(self.cmd, shell=True, stderr=STDOUT, stdout=PIPE)
self.socks = {}
self._output_buffer = []
self.waitForOutput('Starting BsonNetwork v')
time.sleep(1.0)
class BsonNetworkProcess(object):
'''Utility class in order to test BsoNetwork processes.'''
def __init__(self, cmd, **kwargs):
'''Initialize with the given command and arguments.'''
self._parse_arguments(cmd, **kwargs)
self.proc = None
def _parse_arguments(self, cmd, port=None, clientid=None, **kwargs):
'''Parse all the arguments and add them to the command.'''
if isinstance(cmd, str):
cmd = cmd.split(' ')
if '-p' not in cmd and '--port' not in cmd:
if port is None:
port = process.randomPort()
cmd.append('-p')
cmd.append(str(port))
self.port = int(cmd[cmd.index('-p') + 1])
if '-i' not in cmd and '--client-id' not in cmd:
if clientid is None:
clientid = 'server'
cmd.append('-i')
cmd.append(clientid)
self.clientid = cmd[cmd.index('-i') + 1]
self.cmd = ' '.join(cmd)
def start(self):
'''Launch and wait for the process to initialize.'''
print 'starting', self.cmd
self.proc = Popen(self.cmd, shell=True, stderr=STDOUT, stdout=PIPE)
self.socks = {}
self._output_buffer = []
self.waitForOutput('Starting BsonNetwork v')
time.sleep(1.0)
def stop(self):
'''Stop the process, and print out all remaining output.'''
for clientid in self.socks.keys():
self.disconnect(clientid)
self.socks = {}
self.proc.terminate()
if self.proc.poll() is None:
self.proc.kill()
self.proc = None
def __enter__(self):
'''This object can be used in a with clause. starts the process.'''
self.start()
return self
def __exit__(self, type, value, traceback):
'''This object can be used in a with clause. stops the process.'''
self.stop()
def waitForOutput(self, output):
'''Waits for `output` to appear in the stdout or stderr of the process.'''
print '=====> Waiting for: %s' % output
# check backlog first.
for line in self._output_buffer:
if output in line:
self._output_buffer.remove(line) # found it! remove it.
return
signal.signal(signal.SIGALRM, Alarm.handler)
signal.alarm(10)
tic = datetime.datetime.now()
print 'TIC:', tic
try:
line = ''
while output not in line:
if len(line) > 1:
self._output_buffer.append(line)
line = self.proc.stdout.readline().strip()
if len(line) > 1:
print line
except Alarm:
toc = datetime.datetime.now()
print 'RAISE TOC:', toc, ' :: ', (toc - tic)
raise OutputTimeout('Timed out waiting for output: %s' % output)
toc = datetime.datetime.now()
print 'TOC:', toc, ' :: ', (toc - tic)
signal.alarm(0)
def _sendobj(self, clientid, doc):
'''Sends bson document `doc` via socket with `clientid`.'''
self.socks[clientid].send(bson.dumps(doc))
def _recvobj(self, clientid, doc):
'''Receives bson document `doc` via socket `clientid`.'''
print len(bson.dumps(doc))
res = self.socks[clientid].recvobj()
#res = self.socks[clientid].recv(len(bson.dumps(doc)))
print res
return res
def send(self, docs):
'''Sends a collection of bson docs (with _src and _dst)'''
print 'SENDING', docs
if not isinstance(docs, list):
docs = [docs]
for doc in docs:
self._sendobj(doc['_src'], doc)
def receive(self, docs, printall=False):
'''Receives a collection of bson docs (with _src and _dst)'''
if not isinstance(docs, list):
docs = [docs]
for doc in docs:
self.waitForOutput('[%(_dst)s] sending document' % doc)
while printall:
print self.proc.stdout.readline().strip()
doc2 = self._recvobj(doc['_dst'], doc)
if not dicts_equal(doc, doc2):
raise DocumentMismatchError('Document mismatch:\n%s\n%s' % (doc, doc2))
def send_and_receive(self, src, dst, docs):
'''Sends and Receives a collection of bson docs (with _src and _dst)'''
if not isinstance(docs, list):
docs = [docs]
for doc in docs:
doc['_src'] = src
doc['_dst'] = dst
self.send(doc)
self.receive(doc)
def connect(self, clientid, trigger=True):
'''Connects a socket with `clientid` to the process server.'''
sock = None
if trigger:
print '=====> Connecting', clientid, 'to', self.port
sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
sock.connect(('', self.port))
self.socks[clientid] = sock
self.waitForOutput('connection made')
return sock
def disconnect(self, clientid, trigger=True):
'''Disconnects the socket with `clientid` from the process server.'''
if trigger:
sock = self.socks[clientid]
sock.close()
del self.socks[clientid]
self.waitForOutput('[%s] connection closed' % clientid)
def identify(self, clientid, trigger=True):
'''Identifies the socket with `clientid` with the process server.'''
self.waitForOutput('sending identification message')
if trigger:
self._recvobj(clientid, { '_src' : self.clientid } )
self._sendobj(clientid, { '_src' : clientid } )
self.waitForOutput('[%s] connection identified' % clientid)
def reidentify(self, oldid, newid, trigger=True):
'''Re-identifies the socket with `oldid` with `newid`.'''
if trigger:
self.socks[newid] = self.socks[oldid]
del self.socks[oldid]
self._sendobj(newid, { '_src' : newid } )
self.waitForOutput('[%s] connection identified' % newid)
