functional.py

importasyncio
importasyncio.events
importcontextlib
importos
importpprint
importselect
importsocket
importtempfile
importthreading


classFunctionalTestCaseMixin:

defnew_loop(self):
returnasyncio.new_event_loop()

defrun_loop_briefly(self, *, delay=0.01):
self.loop.run_until_complete(asyncio.sleep(delay, loop=self.loop))

defloop_exception_handler(self, loop, context):
self.__unhandled_exceptions.append(context)
self.loop.default_exception_handler(context)

defsetUp(self):
self.loop=self.new_loop()
asyncio.set_event_loop(None)

self.loop.set_exception_handler(self.loop_exception_handler)
self.__unhandled_exceptions= []

# Disable `_get_running_loop`.
self._old_get_running_loop=asyncio.events._get_running_loop
asyncio.events._get_running_loop=lambda: None

deftearDown(self):
try:
self.loop.close()

ifself.__unhandled_exceptions:
print('Unexpected calls to loop.call_exception_handler():')
pprint.pprint(self.__unhandled_exceptions)
self.fail('unexpected calls to loop.call_exception_handler()')

finally:
asyncio.events._get_running_loop=self._old_get_running_loop
asyncio.set_event_loop(None)
self.loop=None

deftcp_server(self, server_prog, *,
family=socket.AF_INET,
addr=None,
timeout=5,
backlog=1,
max_clients=10):

ifaddrisNone:
ifhasattr(socket, 'AF_UNIX') andfamily==socket.AF_UNIX:
withtempfile.NamedTemporaryFile() astmp:
addr=tmp.name
else:
addr= ('127.0.0.1', 0)

sock=socket.socket(family, socket.SOCK_STREAM)

iftimeoutisNone:
raiseRuntimeError('timeout is required')
iftimeout<=0:
raiseRuntimeError('only blocking sockets are supported')
sock.settimeout(timeout)

try:
sock.bind(addr)
sock.listen(backlog)
exceptOSErrorasex:
sock.close()
raiseex

returnTestThreadedServer(
self, sock, server_prog, timeout, max_clients)

deftcp_client(self, client_prog,
family=socket.AF_INET,
timeout=10):

sock=socket.socket(family, socket.SOCK_STREAM)

iftimeoutisNone:
raiseRuntimeError('timeout is required')
iftimeout<=0:
raiseRuntimeError('only blocking sockets are supported')
sock.settimeout(timeout)

returnTestThreadedClient(
self, sock, client_prog, timeout)

defunix_server(self, *args, **kwargs):
ifnothasattr(socket, 'AF_UNIX'):
raiseNotImplementedError
returnself.tcp_server(*args, family=socket.AF_UNIX, **kwargs)

defunix_client(self, *args, **kwargs):
ifnothasattr(socket, 'AF_UNIX'):
raiseNotImplementedError
returnself.tcp_client(*args, family=socket.AF_UNIX, **kwargs)

@contextlib.contextmanager
defunix_sock_name(self):
withtempfile.TemporaryDirectory() astd:
fn=os.path.join(td, 'sock')
try:
yieldfn
finally:
try:
os.unlink(fn)
exceptOSError:
pass

def_abort_socket_test(self, ex):
try:
self.loop.stop()
finally:
self.fail(ex)


##############################################################################
# Socket Testing Utilities
##############################################################################


classTestSocketWrapper:

def__init__(self, sock):
self.__sock=sock

defrecv_all(self, n):
buf=b''
whilelen(buf) <n:
data=self.recv(n-len(buf))
ifdata==b'':
raiseConnectionAbortedError
buf+=data
returnbuf

defstart_tls(self, ssl_context, *,
server_side=False,
server_hostname=None):

ssl_sock=ssl_context.wrap_socket(
self.__sock, server_side=server_side,
server_hostname=server_hostname,
do_handshake_on_connect=False)

try:
ssl_sock.do_handshake()
except:
ssl_sock.close()
raise
finally:
self.__sock.close()

self.__sock=ssl_sock

def__getattr__(self, name):
returngetattr(self.__sock, name)

def__repr__(self):
return'<{} {!r}>'.format(type(self).__name__, self.__sock)


classSocketThread(threading.Thread):

defstop(self):
self._active=False
self.join()

def__enter__(self):
self.start()
returnself

def__exit__(self, *exc):
self.stop()


classTestThreadedClient(SocketThread):

def__init__(self, test, sock, prog, timeout):
threading.Thread.__init__(self, None, None, 'test-client')
self.daemon=True

self._timeout=timeout
self._sock=sock
self._active=True
self._prog=prog
self._test=test

defrun(self):
try:
self._prog(TestSocketWrapper(self._sock))
exceptExceptionasex:
self._test._abort_socket_test(ex)


classTestThreadedServer(SocketThread):

def__init__(self, test, sock, prog, timeout, max_clients):
threading.Thread.__init__(self, None, None, 'test-server')
self.daemon=True

self._clients=0
self._finished_clients=0
self._max_clients=max_clients
self._timeout=timeout
self._sock=sock
self._active=True

self._prog=prog

self._s1, self._s2=socket.socketpair()
self._s1.setblocking(False)

self._test=test

defstop(self):
try:
ifself._s2andself._s2.fileno() !=-1:
try:
self._s2.send(b'stop')
exceptOSError:
pass
finally:
super().stop()

defrun(self):
try:
withself._sock:
self._sock.setblocking(0)
self._run()
finally:
self._s1.close()
self._s2.close()

def_run(self):
whileself._active:
ifself._clients>=self._max_clients:
return

r, w, x=select.select(
 [self._sock, self._s1], [], [], self._timeout)

ifself._s1inr:
return

ifself._sockinr:
try:
conn, addr=self._sock.accept()
exceptBlockingIOError:
continue
exceptsocket.timeout:
ifnotself._active:
return
else:
raise
else:
self._clients+=1
conn.settimeout(self._timeout)
try:
withconn:
self._handle_client(conn)
exceptExceptionasex:
self._active=False
try:
raise
finally:
self._test._abort_socket_test(ex)

def_handle_client(self, sock):
self._prog(TestSocketWrapper(sock))

@property
defaddr(self):
returnself._sock.getsockname()