diff --git a/requirements-dev.txt b/requirements-dev.txt index 6edee41781..1058ffb0ee 100644 --- a/requirements-dev.txt +++ b/requirements-dev.txt @@ -5,6 +5,7 @@ pytest-httpbin==1.0.0 pytest-mock==2.0.0 httpbin==0.7.0 trustme +Werkzeug wheel # Flask Stack diff --git a/tests/test_lowlevel.py b/tests/test_lowlevel.py index 859d07e8a5..2079b4d62e 100644 --- a/tests/test_lowlevel.py +++ b/tests/test_lowlevel.py @@ -2,6 +2,7 @@ import pytest from tests.testserver.server import Server, consume_socket_content +from tests.testserver.werkzeug_server import WerkzeugServer import requests from requests.compat import JSONDecodeError @@ -23,17 +24,16 @@ def echo_response_handler(sock): def test_chunked_upload(): """can safely send generators""" - close_server = threading.Event() - server = Server.basic_response_server(wait_to_close_event=close_server) - data = iter([b"a", b"b", b"c"]) + server = WerkzeugServer.echo_server() + data = iter([b'a', b'b', b'c']) with server as (host, port): - url = f"http://{host}:{port}/" + url = f'http://{host}:{port}/' r = requests.post(url, data=data, stream=True) - close_server.set() # release server block + assert r.content == b'abc' assert r.status_code == 200 - assert r.request.headers["Transfer-Encoding"] == "chunked" + assert r.request.headers['Transfer-Encoding'] == 'chunked' def test_chunked_encoding_error(): diff --git a/tests/testserver/werkzeug_server.py b/tests/testserver/werkzeug_server.py new file mode 100644 index 0000000000..ef8903a675 --- /dev/null +++ b/tests/testserver/werkzeug_server.py @@ -0,0 +1,55 @@ +import multiprocessing +import socket +from contextlib import closing + +from werkzeug import Request, Response, run_simple + + +@Request.application +def echo_application(request): + return Response(request.get_data(), 200, content_type=request.content_type) + + +class WerkzeugServer: + """Realistic WSGI server for unit testing.""" + + SOCKET_CONNECT_TIMEOUT = 2 + + def __init__(self, application, host="localhost", port=0): + super().__init__() + + self.host = host + self.port = port + + # Werkzeug will not automatically pick a valid port for us. + if not self.port: + with closing(socket.socket()) as sock: + sock.bind((self.host, self.port)) + self.port = sock.getsockname()[1] + + self.process = multiprocessing.Process( + target=run_simple, args=(self.host, self.port, application) + ) + + @classmethod + def echo_server(cls): + return WerkzeugServer(echo_application) + + def _socket_is_ready(self): + with closing(socket.socket()) as sock: + sock.settimeout(self.SOCKET_CONNECT_TIMEOUT) + return sock.connect_ex((self.host, self.port)) == 0 + + def __enter__(self): + self.process.start() + + # Confirm that we can actually connect to the socket before we return. + # This protects from flaky tests should the process come up too late. + while not self._socket_is_ready(): + pass + + return self.host, self.port + + def __exit__(self, exc_type, exc_value, traceback): + self.process.terminate() + return False