test_buffered_proto.py

importasyncio
importunittest

fromtest.test_asyncioimportfunctionalasfunc_tests


classReceiveStuffProto(asyncio.BufferedProtocol):
def__init__(self, cb, con_lost_fut):
self.cb=cb
self.con_lost_fut=con_lost_fut

defget_buffer(self, sizehint):
self.buffer=bytearray(100)
returnself.buffer

defbuffer_updated(self, nbytes):
self.cb(self.buffer[:nbytes])

defconnection_lost(self, exc):
ifexcisNone:
self.con_lost_fut.set_result(None)
else:
self.con_lost_fut.set_exception(exc)


classBaseTestBufferedProtocol(func_tests.FunctionalTestCaseMixin):

defnew_loop(self):
raiseNotImplementedError

deftest_buffered_proto_create_connection(self):

NOISE=b'12345678+'*1024

asyncdefclient(addr):
data=b''

defon_buf(buf):
nonlocaldata
data+=buf
ifdata==NOISE:
tr.write(b'1')

conn_lost_fut=self.loop.create_future()

tr, pr=awaitself.loop.create_connection(
lambda: ReceiveStuffProto(on_buf, conn_lost_fut), *addr)

awaitconn_lost_fut

asyncdefon_server_client(reader, writer):
writer.write(NOISE)
awaitreader.readexactly(1)
writer.close()
awaitwriter.wait_closed()

srv=self.loop.run_until_complete(
asyncio.start_server(
on_server_client, '127.0.0.1', 0))

addr=srv.sockets[0].getsockname()
self.loop.run_until_complete(
asyncio.wait_for(client(addr), 5, loop=self.loop))

srv.close()
self.loop.run_until_complete(srv.wait_closed())


classBufferedProtocolSelectorTests(BaseTestBufferedProtocol,
unittest.TestCase):

defnew_loop(self):
returnasyncio.SelectorEventLoop()


@unittest.skipUnless(hasattr(asyncio, 'ProactorEventLoop'), 'Windows only')
classBufferedProtocolProactorTests(BaseTestBufferedProtocol,
unittest.TestCase):

defnew_loop(self):
returnasyncio.ProactorEventLoop()


if__name__=='__main__':
unittest.main()