atst/tests/acceptance/live_server.py
2018-10-15 09:10:08 -04:00

110 lines
3.5 KiB
Python

import gc
import multiprocessing
import socket
import socketserver
import time
from urllib.parse import urlparse, urljoin
# This is adapted from flask-testing, https://github.com/jarus/flask-testing
# Inspired by https://docs.djangoproject.com/en/dev/topics/testing/#django.test.LiveServerTestCase
class LiveServer:
def __init__(self, app, port=5000, timeout=5):
self.app = app
self._configured_port = port
self._timeout = timeout
self._port_value = multiprocessing.Value("i", self._configured_port)
@property
def server_url(self):
return "http://localhost:%s" % self._port_value.value
def spawn_live_server(self):
self._process = None
port_value = self._port_value
def worker(app, port):
# Based on solution: http://stackoverflow.com/a/27598916
# Monkey-patch the server_bind so we can determine the port bound
# by Flask. This handles the case where the port specified is `0`,
# which means that the OS chooses the port. This is the only known
# way (currently) of getting the port out of Flask once we call
# `run`.
original_socket_bind = socketserver.TCPServer.server_bind
def socket_bind_wrapper(self):
ret = original_socket_bind(self)
# Get the port and save it into the port_value, so the parent
# process can read it.
(_, port) = self.socket.getsockname()
port_value.value = port
socketserver.TCPServer.server_bind = original_socket_bind
return ret
socketserver.TCPServer.server_bind = socket_bind_wrapper
app.run(port=port, use_reloader=False)
self._process = multiprocessing.Process(
target=worker, args=(self.app, self._configured_port)
)
self._process.start()
# We must wait for the server to start listening, but give up
# after a specified maximum timeout
start_time = time.time()
while True:
elapsed_time = time.time() - start_time
if elapsed_time > self._timeout:
raise RuntimeError(
"Failed to start the server after %d seconds. " % self._timeout
)
if self._can_ping_server():
break
def _can_ping_server(self):
host, port = self.address
if port == 0:
# Port specified by the user was 0, and the OS has not yet assigned
# the proper port.
return False
sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
try:
sock.connect((host, port))
except socket.error as e:
success = False
else:
success = True
finally:
sock.close()
return success
@property
def address(self):
"""
Gets the server address used to test the connection with a socket.
Respects both the LIVESERVER_PORT config value and overriding server_url
"""
parts = urlparse(self.server_url)
host = parts.hostname
port = parts.port
if port is None:
if parts.scheme == "http":
port = 80
elif parts.scheme == "https":
port = 443
else:
raise RuntimeError("Unsupported server url scheme: %s" % parts.scheme)
return host, port
def terminate(self):
if self._process:
self._process.terminate()