From d119ab948352958e5b7de1b9b27a7bffd95cc7dd Mon Sep 17 00:00:00 2001
From: Josh Boyer <jwboyer@redhat.com>
Date: Mon, 15 May 2023 07:26:54 -0400
Subject: [PATCH 6/6] Add Vary cookie tests

---
 flask/testsuite/basic.py | 97 +++++++++++++++++++++++++++++++++++++++-
 1 file changed, 95 insertions(+), 2 deletions(-)

diff --git a/flask/testsuite/basic.py b/flask/testsuite/basic.py
index 0fffd4ee..eb5d65a8 100644
--- a/flask/testsuite/basic.py
+++ b/flask/testsuite/basic.py
@@ -18,9 +18,10 @@ from datetime import datetime
 from threading import Thread
 from flask.testsuite import FlaskTestCase, emits_module_deprecation_warning
 from flask._compat import text_type
-from werkzeug.exceptions import BadRequest, NotFound
+import werkzeug.serving
+from werkzeug.exceptions import BadRequest, NotFound, Forbidden
 from werkzeug.http import parse_date
-from werkzeug.routing import BuildError
+from werkzeug.routing import BuildError, RequestRedirect
 
 
 class BasicFunctionalityTestCase(FlaskTestCase):
@@ -388,6 +389,98 @@ class BasicFunctionalityTestCase(FlaskTestCase):
         app.config['SESSION_REFRESH_EACH_REQUEST'] = False
         run_test(expect_header=False)
 
+    def test_session_vary_cookie(self):
+        app = flask.Flask(__name__)
+        app.testing = True
+        app.secret_key = 'dev key'
+        is_permanent = True
+
+        @app.route("/set")
+        def set_session():
+            flask.session["test"] = "test"
+            return ""
+
+        @app.route("/get")
+        def get():
+            return flask.session.get("test")
+
+        @app.route("/getitem")
+        def getitem():
+            return flask.session["test"]
+
+        @app.route("/setdefault")
+        def setdefault():
+            return flask.session.setdefault("test", "default")
+
+        @app.route("/clear")
+        def clear():
+            flask.session.clear()
+            return ""
+
+        @app.route("/vary-cookie-header-set")
+        def vary_cookie_header_set():
+            response = flask.Response()
+            response.vary.add("Cookie")
+            flask.session["test"] = "test"
+            return response
+
+        @app.route("/vary-header-set")
+        def vary_header_set():
+            response = flask.Response()
+            response.vary.update(("Accept-Encoding", "Accept-Language"))
+            flask.session["test"] = "test"
+            return response
+
+        @app.route("/no-vary-header")
+        def no_vary_header():
+            return ""
+
+        c = app.test_client()
+
+        def expect(path, header_value="Cookie"):
+            rv = c.get(path)
+
+            if header_value:
+                #print "%s" % rv.headers.get_all("Vary")
+                # The 'Vary' key should exist in the headers only once.
+                assert len(rv.headers.get_all("Vary")) == 1
+                assert rv.headers["Vary"] == header_value
+            else:
+                assert "Vary" not in rv.headers
+
+        expect("/set")
+        expect("/get")
+        expect("/getitem")
+        expect("/setdefault")
+        expect("/clear")
+        expect("/vary-cookie-header-set")
+        expect("/vary-header-set", "Accept-Encoding, Accept-Language, Cookie")
+        expect("/no-vary-header", None)
+
+    def test_session_refresh_vary(self):
+        app = flask.Flask(__name__)
+        app.testing = True
+        app.secret_key = 'key dev'
+        is_permanent = True
+
+        @app.route("/login")
+        def login():
+            flask.session["user_id"] = 1
+            flask.session.permanent = True
+            return ""
+
+        @app.route("/ignored")
+        def ignored():
+            return ""
+
+        c = app.test_client()
+
+        rv = c.get("/login")
+        assert rv.headers["Vary"] == "Cookie"
+        rv = c.get("/ignored")
+        assert rv.headers["Vary"] == "Cookie"
+
+
     def test_flashes(self):
         app = flask.Flask(__name__)
         app.secret_key = 'testkey'
-- 
2.31.1