def setUp(self):
self.conns = [None] * 2
self.server = None
self.loop = get_loop()
def serve(reader, writer):
self.conns[1] = Connection(self.loop, reader, writer)
@asyncio.coroutine
def connect():
host, port = '127.0.0.10', 5000
# let a server listen
self.server = (yield from asyncio.start_server(serve, host, port, loop=self.loop))
# connect a client
reader, writer = yield from asyncio.open_connection(host, port, loop=self.loop)
self.conns[0] = Connection(self.loop, reader, writer)
def close():
if self.server:
self.server.close()
for conn in self.conns:
if conn:
run_coroutine_threadsafe(conn.close(), loop=self.loop).result()
self.addCleanup(close)
run_coroutine_threadsafe(connect(), loop=self.loop).result()
def test_sendfile(self):
sendsock, recvsock = socket.socketpair()
src = tempfile.NamedTemporaryFile(prefix='bndl-sendfiletest-')
src.file.write(self.data)
received = []
def receive():
while sum(len(buffer) for buffer in received) < len(self.data):
buffer = recvsock.recv(len(self.data))
received.append(buffer)
receiver = threading.Thread(target=receive)
receiver.start()
loop = get_loop()
run_coroutine_threadsafe(sendfile(sendsock.fileno(),
src.file.fileno(),
0,
len(self.data),
loop=loop),
loop=loop).result()
receiver.join()
received = b''.join(received)
self.assertEqual(self.data, received)
def stop_async(self):
if not self.loop.is_closed():
return aio.run_coroutine_threadsafe(self.stop(), self.loop)
else:
self._stop_tasks()
fut = concurrent.futures.Future()
fut.set_result(None)
return fut
def disconnect_async(self, reason='', active=True):
if not self.loop.is_closed():
return aio.run_coroutine_threadsafe(
self.disconnect(reason, active), self.loop)
else:
self._stop_tasks()
fut = concurrent.futures.Future()
fut.set_result(None)
return fut
def stop_nodes(nodes):
stops = [
run_coroutine_threadsafe(node.stop(), node.loop) for node in nodes
]
for stop in stops:
try:
stop.result()
except Exception:
logger.exception('Unable to stop node')
def start_nodes(nodes):
started = []
starts = [
run_coroutine_threadsafe(node.start(), node.loop) for node in nodes
]
try:
for start in starts:
start.result()
except Exception:
logger.exception('Unable to start node')
stop_nodes(started)
raise
def test_send_on_disk(self):
in_memory = StorageContainerFactory('memory', 'pickle')('a')
on_disk = StorageContainerFactory('disk', 'pickle')('b')
to_disk = StorageContainerFactory('memory', 'pickle')('c')
in_memory.write(self.data[0::3])
on_disk.write(self.data[1::3])
to_disk.write(self.data[2::3])
to_disk.to_disk()
def connected(reader, writer):
conn = Connection(self.loop, reader, writer)
yield from conn.send(in_memory)
yield from conn.send(on_disk)
yield from conn.send(to_disk)
@asyncio.coroutine
def run_pair():
server = yield from asyncio.start_server(connected,
'0.0.0.0',
0,
loop=self.loop)
socket = server.sockets[0]
host, port = socket.getsockname()[:2]
reader, writer = yield from asyncio.open_connection(host,
port,
loop=self.loop)
conn = Connection(self.loop, reader, writer)
a = yield from conn.recv()
b = yield from conn.recv()
c = yield from conn.recv()
self.assertEqual(in_memory.read(), a.read())
self.assertEqual(on_disk.read(), b.read())
self.assertEqual(to_disk.read(), c.read())
run_coroutine_threadsafe(run_pair(), self.loop).result()
def test_copyfile(self):
src = tempfile.NamedTemporaryFile(prefix='bndl-sendfiletest-')
src.file.write(self.data)
src.file.seek(0)
size = os.stat(src.fileno()).st_size
dst = tempfile.NamedTemporaryFile(prefix='bndl-sendfiletest-')
loop = get_loop()
run_coroutine_threadsafe(sendfile(dst.file.fileno(),
src.file.fileno(),
0,
size,
loop=loop),
loop=loop).result()
self.assertEqual(size, os.stat(dst.fileno()).st_size)
src.file.seek(0)
dst.file.seek(0)
self.assertEqual(src.file.read(), dst.file.read())
src.close()
dst.close()
def recv(self, conn, timeout=None):
@asyncio.coroutine
def _recv():
return (yield from conn.recv(timeout))
return run_coroutine_threadsafe(_recv(), loop=self.loop).result()
def send(self, conn, msg):
run_coroutine_threadsafe(conn.send(msg, drain=False), loop=self.loop).result()
def close():
if self.server:
self.server.close()
for conn in self.conns:
if conn:
run_coroutine_threadsafe(conn.close(), loop=self.loop).result()
def start_async(self):
return aio.run_coroutine_threadsafe(self.start(), self.loop)
def run_coro(self, coro):
return run_coroutine_threadsafe(coro, loop=self.loop)
def __call__(self, *args, **kwargs):
request = self._request(*args, **kwargs)
return run_coroutine_threadsafe(request, self.peer.loop)