From 04c21387dbc454fb59cae9c9bb21a13358cb5a34 Mon Sep 17 00:00:00 2001 From: David Lord Date: Wed, 12 Apr 2023 10:33:52 -0700 Subject: [PATCH] update test cookie handling for Werkzeug 2.3 --- src/flask/testing.py | 66 +++++++++++++++++++++++-------------------- tests/test_basic.py | 16 +++++++---- tests/test_json.py | 19 ------------- tests/test_testing.py | 4 +-- 4 files changed, 47 insertions(+), 58 deletions(-) diff --git a/src/flask/testing.py b/src/flask/testing.py index 8cb2d1bd..a972a3f5 100644 --- a/src/flask/testing.py +++ b/src/flask/testing.py @@ -11,7 +11,6 @@ from werkzeug.test import Client from werkzeug.wrappers import Request as BaseRequest from .cli import ScriptInfo -from .globals import _cv_request from .sessions import SessionMixin if t.TYPE_CHECKING: # pragma: no cover @@ -137,40 +136,45 @@ class FlaskClient(Client): :meth:`~flask.Flask.test_request_context` which are directly passed through. """ - if self.cookie_jar is None: - raise RuntimeError( - "Session transactions only make sense with cookies enabled." + # new cookie interface for Werkzeug >= 2.3 + cookie_storage = self._cookies if hasattr(self, "_cookies") else self.cookie_jar + + if cookie_storage is None: + raise TypeError( + "Cookies are disabled. Create a client with 'use_cookies=True'." ) + app = self.application - environ_overrides = kwargs.setdefault("environ_overrides", {}) - self.cookie_jar.inject_wsgi(environ_overrides) - outer_reqctx = _cv_request.get(None) - with app.test_request_context(*args, **kwargs) as c: - session_interface = app.session_interface - sess = session_interface.open_session(app, c.request) - if sess is None: - raise RuntimeError( - "Session backend did not open a session. Check the configuration" - ) + ctx = app.test_request_context(*args, **kwargs) - # Since we have to open a new request context for the session - # handling we want to make sure that we hide out own context - # from the caller. By pushing the original request context - # (or None) on top of this and popping it we get exactly that - # behavior. It's important to not use the push and pop - # methods of the actual request context object since that would - # mean that cleanup handlers are called - token = _cv_request.set(outer_reqctx) # type: ignore[arg-type] - try: - yield sess - finally: - _cv_request.reset(token) + if hasattr(self, "_add_cookies_to_wsgi"): + self._add_cookies_to_wsgi(ctx.request.environ) + else: + self.cookie_jar.inject_wsgi(ctx.request.environ) # type: ignore[union-attr] - resp = app.response_class() - if not session_interface.is_null_session(sess): - session_interface.save_session(app, sess, resp) - headers = resp.get_wsgi_headers(c.request.environ) - self.cookie_jar.extract_wsgi(c.request.environ, headers) + with ctx: + sess = app.session_interface.open_session(app, ctx.request) + + if sess is None: + raise RuntimeError("Session backend did not open a session.") + + yield sess + resp = app.response_class() + + if app.session_interface.is_null_session(sess): + return + + with ctx: + app.session_interface.save_session(app, sess, resp) + + if hasattr(self, "_update_cookies_from_response"): + self._update_cookies_from_response( + ctx.request.host.partition(":")[0], resp.headers.getlist("Set-Cookie") + ) + else: + self.cookie_jar.extract_wsgi( # type: ignore[union-attr] + ctx.request.environ, resp.headers + ) def _copy_environ(self, other): out = {**self.environ_base, **other} diff --git a/tests/test_basic.py b/tests/test_basic.py index 9aca6679..a622fa93 100644 --- a/tests/test_basic.py +++ b/tests/test_basic.py @@ -261,8 +261,9 @@ def test_session_using_server_name(app, client): return "Hello World" rv = client.get("/", "http://example.com/") - assert "domain=.example.com" in rv.headers["set-cookie"].lower() - assert "httponly" in rv.headers["set-cookie"].lower() + cookie = rv.headers["set-cookie"].lower() + # or condition for Werkzeug < 2.3 + assert "domain=example.com" in cookie or "domain=.example.com" in cookie def test_session_using_server_name_and_port(app, client): @@ -274,8 +275,9 @@ def test_session_using_server_name_and_port(app, client): return "Hello World" rv = client.get("/", "http://example.com:8080/") - assert "domain=.example.com" in rv.headers["set-cookie"].lower() - assert "httponly" in rv.headers["set-cookie"].lower() + cookie = rv.headers["set-cookie"].lower() + # or condition for Werkzeug < 2.3 + assert "domain=example.com" in cookie or "domain=.example.com" in cookie def test_session_using_server_name_port_and_path(app, client): @@ -337,7 +339,8 @@ def test_session_using_session_settings(app, client): rv = client.get("/", "http://www.example.com:8080/test/") cookie = rv.headers["set-cookie"].lower() - assert "domain=.example.com" in cookie + # or condition for Werkzeug < 2.3 + assert "domain=example.com" in cookie or "domain=.example.com" in cookie assert "path=/" in cookie assert "secure" in cookie assert "httponly" not in cookie @@ -346,7 +349,8 @@ def test_session_using_session_settings(app, client): rv = client.get("/clear", "http://www.example.com:8080/test/") cookie = rv.headers["set-cookie"].lower() assert "session=;" in cookie - assert "domain=.example.com" in cookie + # or condition for Werkzeug < 2.3 + assert "domain=example.com" in cookie or "domain=.example.com" in cookie assert "path=/" in cookie assert "secure" in cookie assert "samesite" in cookie diff --git a/tests/test_json.py b/tests/test_json.py index ba9f38dc..500eb64d 100644 --- a/tests/test_json.py +++ b/tests/test_json.py @@ -267,25 +267,6 @@ def _has_encoding(name): return False -@pytest.mark.skipif( - not _has_encoding("euc-kr"), reason="The euc-kr encoding is required." -) -def test_modified_url_encoding(app, client): - class ModifiedRequest(flask.Request): - url_charset = "euc-kr" - - app.request_class = ModifiedRequest - app.url_map.charset = "euc-kr" - - @app.route("/") - def index(): - return flask.request.args["foo"] - - rv = client.get("/", query_string={"foo": "정상처리"}, charset="euc-kr") - assert rv.status_code == 200 - assert rv.get_data(as_text=True) == "정상처리" - - def test_json_key_sorting(app, client): app.debug = True assert app.json.sort_keys diff --git a/tests/test_testing.py b/tests/test_testing.py index 40f6fe1b..14c5ade0 100644 --- a/tests/test_testing.py +++ b/tests/test_testing.py @@ -206,10 +206,10 @@ def test_session_transactions_keep_context(app, client, req_ctx): def test_session_transaction_needs_cookies(app): c = app.test_client(use_cookies=False) - with pytest.raises(RuntimeError) as e: + + with pytest.raises(TypeError, match="Cookies are disabled."): with c.session_transaction(): pass - assert "cookies" in str(e.value) def test_test_client_context_binding(app, client):