fixed issue 3396

This commit is contained in:
Rohan-Salwan 2021-09-10 20:35:10 +05:30
parent 313a70da6e
commit d589f2d5f7
2 changed files with 16 additions and 9 deletions

View file

@ -21,7 +21,7 @@ Unreleased
:issue:`4096` :issue:`4096`
- The CLI loader handles ``**kwargs`` in a ``create_app`` function. - The CLI loader handles ``**kwargs`` in a ``create_app`` function.
:issue:`4170` :issue:`4170`
- Fix Unable to extend FlaskClient with follow_redirects. :issue:`3396`
Version 2.0.1 Version 2.0.1
------------- -------------

View file

@ -108,6 +108,7 @@ class FlaskClient(Client):
application: "Flask" application: "Flask"
preserve_context = False preserve_context = False
arg = None
def __init__(self, *args: t.Any, **kwargs: t.Any) -> None: def __init__(self, *args: t.Any, **kwargs: t.Any) -> None:
super().__init__(*args, **kwargs) super().__init__(*args, **kwargs)
@ -189,9 +190,11 @@ class FlaskClient(Client):
"flask._preserve_context": self.preserve_context, "flask._preserve_context": self.preserve_context,
} }
if len(args) < 1:
pass
else:
self.arg = arg = args[0]
if not kwargs and len(args) == 1: if not kwargs and len(args) == 1:
arg = args[0]
if isinstance(arg, werkzeug.test.EnvironBuilder): if isinstance(arg, werkzeug.test.EnvironBuilder):
builder = copy(arg) builder = copy(arg)
builder.environ_base = copy_environ(builder.environ_base or {}) builder.environ_base = copy_environ(builder.environ_base or {})
@ -205,13 +208,17 @@ class FlaskClient(Client):
request.environ = copy_environ(request.environ) request.environ = copy_environ(request.environ)
if request is None: if request is None:
kwargs["environ_base"] = copy_environ(kwargs.get("environ_base", {})) if isinstance(self.arg, werkzeug.test.EnvironBuilder):
builder = EnvironBuilder(self.application, *args, **kwargs) builder = copy(arg)
builder.environ_base = copy_environ(builder.environ_base or {})
try:
request = builder.get_request() request = builder.get_request()
finally: else:
builder.close() kwargs["environ_base"] = copy_environ(kwargs.get("environ_base", {}))
builder = EnvironBuilder(self.application, *args, **kwargs)
try:
request = builder.get_request()
finally:
builder.close()
return super().open( # type: ignore return super().open( # type: ignore
request, request,