overriding FlaskClient.open works with redirects

This commit is contained in:
David Lord 2022-02-09 11:13:03 -08:00
parent 7c5f17a55e
commit 1be65b1b69
No known key found for this signature in database
GPG Key ID: 7A1C87E3F5BC42A8
2 changed files with 33 additions and 30 deletions

View File

@ -37,6 +37,8 @@ Unreleased
:pr:`4303` :pr:`4303`
- The CLI uses ``importlib.metadata`` instead of ``setuptools`` to - The CLI uses ``importlib.metadata`` instead of ``setuptools`` to
load command entry points. :issue:`4419` load command entry points. :issue:`4419`
- Overriding ``FlaskClient.open`` will not cause an error on redirect.
:issue:`3396`
Version 2.0.3 Version 2.0.3

View File

@ -172,6 +172,22 @@ class FlaskClient(Client):
headers = resp.get_wsgi_headers(c.request.environ) headers = resp.get_wsgi_headers(c.request.environ)
self.cookie_jar.extract_wsgi(c.request.environ, headers) self.cookie_jar.extract_wsgi(c.request.environ, headers)
def _copy_environ(self, other):
return {
**self.environ_base,
**other,
"flask._preserve_context": self.preserve_context,
}
def _request_from_builder_args(self, args, kwargs):
kwargs["environ_base"] = self._copy_environ(kwargs.get("environ_base", {}))
builder = EnvironBuilder(self.application, *args, **kwargs)
try:
return builder.get_request()
finally:
builder.close()
def open( def open(
self, self,
*args: t.Any, *args: t.Any,
@ -179,39 +195,24 @@ class FlaskClient(Client):
follow_redirects: bool = False, follow_redirects: bool = False,
**kwargs: t.Any, **kwargs: t.Any,
) -> "TestResponse": ) -> "TestResponse":
# Same logic as super.open, but apply environ_base and preserve_context. if args and isinstance(
request = None args[0], (werkzeug.test.EnvironBuilder, dict, BaseRequest)
):
def copy_environ(other): if isinstance(args[0], werkzeug.test.EnvironBuilder):
return { builder = copy(args[0])
**self.environ_base, builder.environ_base = self._copy_environ(builder.environ_base or {})
**other,
"flask._preserve_context": self.preserve_context,
}
if not kwargs and len(args) == 1:
arg = args[0]
if isinstance(arg, werkzeug.test.EnvironBuilder):
builder = copy(arg)
builder.environ_base = copy_environ(builder.environ_base or {})
request = builder.get_request() request = builder.get_request()
elif isinstance(arg, dict): elif isinstance(args[0], dict):
request = EnvironBuilder.from_environ( request = EnvironBuilder.from_environ(
arg, app=self.application, environ_base=copy_environ({}) args[0], app=self.application, environ_base=self._copy_environ({})
).get_request() ).get_request()
elif isinstance(arg, BaseRequest): else:
request = copy(arg) # isinstance(args[0], BaseRequest)
request.environ = copy_environ(request.environ) request = copy(args[0])
request.environ = self._copy_environ(request.environ)
if request is None: else:
kwargs["environ_base"] = copy_environ(kwargs.get("environ_base", {})) # request is None
builder = EnvironBuilder(self.application, *args, **kwargs) request = self._request_from_builder_args(args, kwargs)
try:
request = builder.get_request()
finally:
builder.close()
return super().open( return super().open(
request, request,