test_server.py

importasyncio
importos
importsocket
importtime
importthreading
importunittest

fromtest.supportimportsocket_helper
fromtest.test_asyncioimportutilsastest_utils
fromtest.test_asyncioimportfunctionalasfunc_tests


deftearDownModule():
asyncio.set_event_loop_policy(None)


classBaseStartServer(func_tests.FunctionalTestCaseMixin):

defnew_loop(self):
raiseNotImplementedError

deftest_start_server_1(self):
HELLO_MSG=b'1'*1024*5+b'\n'

defclient(sock, addr):
foriinrange(10):
time.sleep(0.2)
ifsrv.is_serving():
break
else:
raiseRuntimeError

sock.settimeout(2)
sock.connect(addr)
sock.send(HELLO_MSG)
sock.recv_all(1)
sock.close()

asyncdefserve(reader, writer):
awaitreader.readline()
main_task.cancel()
writer.write(b'1')
writer.close()
awaitwriter.wait_closed()

asyncdefmain(srv):
asyncwithsrv:
awaitsrv.serve_forever()

srv=self.loop.run_until_complete(asyncio.start_server(
serve, socket_helper.HOSTv4, 0, start_serving=False))

self.assertFalse(srv.is_serving())

main_task=self.loop.create_task(main(srv))

addr=srv.sockets[0].getsockname()
withself.assertRaises(asyncio.CancelledError):
withself.tcp_client(lambdasock: client(sock, addr)):
self.loop.run_until_complete(main_task)

self.assertEqual(srv.sockets, ())

self.assertIsNone(srv._sockets)
self.assertIsNone(srv._waiters)
self.assertFalse(srv.is_serving())

withself.assertRaisesRegex(RuntimeError, r'is closed'):
self.loop.run_until_complete(srv.serve_forever())


classSelectorStartServerTests(BaseStartServer, unittest.TestCase):

defnew_loop(self):
returnasyncio.SelectorEventLoop()

@socket_helper.skip_unless_bind_unix_socket
deftest_start_unix_server_1(self):
HELLO_MSG=b'1'*1024*5+b'\n'
started=threading.Event()

defclient(sock, addr):
sock.settimeout(2)
started.wait(5)
sock.connect(addr)
sock.send(HELLO_MSG)
sock.recv_all(1)
sock.close()

asyncdefserve(reader, writer):
awaitreader.readline()
main_task.cancel()
writer.write(b'1')
writer.close()
awaitwriter.wait_closed()

asyncdefmain(srv):
asyncwithsrv:
self.assertFalse(srv.is_serving())
awaitsrv.start_serving()
self.assertTrue(srv.is_serving())
started.set()
awaitsrv.serve_forever()

withtest_utils.unix_socket_path() asaddr:
srv=self.loop.run_until_complete(asyncio.start_unix_server(
serve, addr, start_serving=False))

main_task=self.loop.create_task(main(srv))

withself.assertRaises(asyncio.CancelledError):
withself.unix_client(lambdasock: client(sock, addr)):
self.loop.run_until_complete(main_task)

self.assertEqual(srv.sockets, ())

self.assertIsNone(srv._sockets)
self.assertIsNone(srv._waiters)
self.assertFalse(srv.is_serving())

withself.assertRaisesRegex(RuntimeError, r'is closed'):
self.loop.run_until_complete(srv.serve_forever())


classTestServer2(unittest.IsolatedAsyncioTestCase):

asyncdeftest_wait_closed_basic(self):
asyncdefserve(rd, wr):
try:
awaitrd.read()
finally:
wr.close()
awaitwr.wait_closed()

srv=awaitasyncio.start_server(serve, socket_helper.HOSTv4, 0)
self.addCleanup(srv.close)

# active count = 0, not closed: should block
task1=asyncio.create_task(srv.wait_closed())
awaitasyncio.sleep(0)
self.assertFalse(task1.done())

# active count != 0, not closed: should block
addr=srv.sockets[0].getsockname()
 (rd, wr) =awaitasyncio.open_connection(addr[0], addr[1])
task2=asyncio.create_task(srv.wait_closed())
awaitasyncio.sleep(0)
self.assertFalse(task1.done())
self.assertFalse(task2.done())

srv.close()
awaitasyncio.sleep(0)
# active count != 0, closed: should block
task3=asyncio.create_task(srv.wait_closed())
awaitasyncio.sleep(0)
self.assertFalse(task1.done())
self.assertFalse(task2.done())
self.assertFalse(task3.done())

wr.close()
awaitwr.wait_closed()
# active count == 0, closed: should unblock
awaittask1
awaittask2
awaittask3
awaitsrv.wait_closed() # Return immediately

asyncdeftest_wait_closed_race(self):
# Test a regression in 3.12.0, should be fixed in 3.12.1
asyncdefserve(rd, wr):
try:
awaitrd.read()
finally:
wr.close()
awaitwr.wait_closed()

srv=awaitasyncio.start_server(serve, socket_helper.HOSTv4, 0)
self.addCleanup(srv.close)

task=asyncio.create_task(srv.wait_closed())
awaitasyncio.sleep(0)
self.assertFalse(task.done())
addr=srv.sockets[0].getsockname()
 (rd, wr) =awaitasyncio.open_connection(addr[0], addr[1])
loop=asyncio.get_running_loop()
loop.call_soon(srv.close)
loop.call_soon(wr.close)
awaitsrv.wait_closed()

asyncdeftest_close_clients(self):
asyncdefserve(rd, wr):
try:
awaitrd.read()
finally:
wr.close()
awaitwr.wait_closed()

srv=awaitasyncio.start_server(serve, socket_helper.HOSTv4, 0)
self.addCleanup(srv.close)

addr=srv.sockets[0].getsockname()
 (rd, wr) =awaitasyncio.open_connection(addr[0], addr[1])
self.addCleanup(wr.close)

task=asyncio.create_task(srv.wait_closed())
awaitasyncio.sleep(0)
self.assertFalse(task.done())

srv.close()
srv.close_clients()
awaitasyncio.sleep(0)
awaitasyncio.sleep(0)
self.assertTrue(task.done())

asyncdeftest_abort_clients(self):
asyncdefserve(rd, wr):
fut.set_result((rd, wr))
awaitwr.wait_closed()

fut=asyncio.Future()
srv=awaitasyncio.start_server(serve, socket_helper.HOSTv4, 0)
self.addCleanup(srv.close)

addr=srv.sockets[0].getsockname()
 (c_rd, c_wr) =awaitasyncio.open_connection(addr[0], addr[1], limit=4096)
self.addCleanup(c_wr.close)

 (s_rd, s_wr) =awaitfut

# Limit the socket buffers so we can more reliably overfill them
s_sock=s_wr.get_extra_info('socket')
s_sock.setsockopt(socket.SOL_SOCKET, socket.SO_SNDBUF, 65536)
c_sock=c_wr.get_extra_info('socket')
c_sock.setsockopt(socket.SOL_SOCKET, socket.SO_RCVBUF, 65536)

# Get the reader in to a paused state by sending more than twice
# the configured limit
s_wr.write(b'a'*4096)
s_wr.write(b'a'*4096)
s_wr.write(b'a'*4096)
whilec_wr.transport.is_reading():
awaitasyncio.sleep(0)

# Get the writer in a waiting state by sending data until the
# kernel stops accepting more data in the send buffer.
# gh-122136: getsockopt() does not reliably report the buffer size
# available for message content.
# We loop until we start filling up the asyncio buffer.
# To avoid an infinite loop we cap at 10 times the expected value
c_bufsize=c_sock.getsockopt(socket.SOL_SOCKET, socket.SO_RCVBUF)
s_bufsize=s_sock.getsockopt(socket.SOL_SOCKET, socket.SO_SNDBUF)
foriinrange(10):
s_wr.write(b'a'*c_bufsize)
s_wr.write(b'a'*s_bufsize)
ifs_wr.transport.get_write_buffer_size() >0:
break
self.assertNotEqual(s_wr.transport.get_write_buffer_size(), 0)

task=asyncio.create_task(srv.wait_closed())
awaitasyncio.sleep(0)
self.assertFalse(task.done())

srv.close()
srv.abort_clients()
awaitasyncio.sleep(0)
awaitasyncio.sleep(0)
self.assertTrue(task.done())


# Test the various corner cases of Unix server socket removal
classUnixServerCleanupTests(unittest.IsolatedAsyncioTestCase):
@socket_helper.skip_unless_bind_unix_socket
asyncdeftest_unix_server_addr_cleanup(self):
# Default scenario
withtest_utils.unix_socket_path() asaddr:
asyncdefserve(*args):
pass

srv=awaitasyncio.start_unix_server(serve, addr)

srv.close()
self.assertFalse(os.path.exists(addr))

@socket_helper.skip_unless_bind_unix_socket
asyncdeftest_unix_server_sock_cleanup(self):
# Using already bound socket
withtest_utils.unix_socket_path() asaddr:
asyncdefserve(*args):
pass

withsocket.socket(socket.AF_UNIX, socket.SOCK_STREAM) assock:
sock.bind(addr)

srv=awaitasyncio.start_unix_server(serve, sock=sock)

srv.close()
self.assertFalse(os.path.exists(addr))

@socket_helper.skip_unless_bind_unix_socket
asyncdeftest_unix_server_cleanup_gone(self):
# Someone else has already cleaned up the socket
withtest_utils.unix_socket_path() asaddr:
asyncdefserve(*args):
pass

withsocket.socket(socket.AF_UNIX, socket.SOCK_STREAM) assock:
sock.bind(addr)

srv=awaitasyncio.start_unix_server(serve, sock=sock)

os.unlink(addr)

srv.close()

@socket_helper.skip_unless_bind_unix_socket
asyncdeftest_unix_server_cleanup_replaced(self):
# Someone else has replaced the socket with their own
withtest_utils.unix_socket_path() asaddr:
asyncdefserve(*args):
pass

srv=awaitasyncio.start_unix_server(serve, addr)

os.unlink(addr)
withsocket.socket(socket.AF_UNIX, socket.SOCK_STREAM) assock:
sock.bind(addr)

srv.close()
self.assertTrue(os.path.exists(addr))

@socket_helper.skip_unless_bind_unix_socket
asyncdeftest_unix_server_cleanup_prevented(self):
# Automatic cleanup explicitly disabled
withtest_utils.unix_socket_path() asaddr:
asyncdefserve(*args):
pass

srv=awaitasyncio.start_unix_server(serve, addr, cleanup_socket=False)

srv.close()
self.assertTrue(os.path.exists(addr))


@unittest.skipUnless(hasattr(asyncio, 'ProactorEventLoop'), 'Windows only')
classProactorStartServerTests(BaseStartServer, unittest.TestCase):

defnew_loop(self):
returnasyncio.ProactorEventLoop()


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