singingwolfboy · September 21, 2024 23:28 · Feb 26, 2016 · Feb 26, 2016
diff --git a/flask_csrf_test_client.py b/flask_csrf_test_client.py
@@ -28,7 +28,7 @@ def set_cookie(self, key, value='', *args, **kwargs):
 
     def delete_cookie(self, key, *args, **kwargs):
         "Delete the cookie on the Flask test client."
-        server_name = flask.current_app.config.get("SERVER_NAME", "localhost")
+        server_name = flask.current_app.config["SERVER_NAME"] or "localhost"
         return self.client.delete_cookie(
             server_name, key=key, *args, **kwargs
         )

diff --git a/flask_csrf_test_client.py b/flask_csrf_test_client.py
@@ -0,0 +1,91 @@
+# Want to run your Flask tests with CSRF protections turned on, to make sure
+# that CSRF works properly in production as well? Here's an excellent way
+# to do it!
+
+# First some imports. I'm assuming you're using Flask-WTF for CSRF protection.
+import flask
+from flask.testing import FlaskClient as BaseFlaskClient
+from flask_wtf.csrf import generate_csrf
+
+# Flask's assumptions about an incoming request don't quite match up with
+# what the test client provides in terms of manipulating cookies, and the
+# CSRF system depends on cookies working correctly. This little class is a
+# fake request that forwards along requests to the test client for setting
+# cookies.
+class RequestShim(object):
+    """
+    A fake request that proxies cookie-related methods to a Flask test client.
+    """
+    def __init__(self, client):
+        self.client = client
+
+    def set_cookie(self, key, value='', *args, **kwargs):
+        "Set the cookie on the Flask test client."
+        server_name = flask.current_app.config["SERVER_NAME"] or "localhost"
+        return self.client.set_cookie(
+            server_name, key=key, value=value, *args, **kwargs
+        )
+
+    def delete_cookie(self, key, *args, **kwargs):
+        "Delete the cookie on the Flask test client."
+        server_name = flask.current_app.config.get("SERVER_NAME", "localhost")
+        return self.client.delete_cookie(
+            server_name, key=key, *args, **kwargs
+        )
+
+# We're going to extend Flask's built-in test client class, so that it knows
+# how to look up CSRF tokens for you!
+class FlaskClient(BaseFlaskClient):
+    @property
+    def csrf_token(self):
+        # First, we'll wrap our request shim around the test client, so that
+        # it will work correctly when Flask asks it to set a cookie.
+        request = RequestShim(self) 
+        # Next, we need to look up any cookies that might already exist on
+        # this test client, such as the secure cookie that powers `flask.session`,
+        # and make a test request context that has those cookies in it.
+        environ_overrides = {}
+        self.cookie_jar.inject_wsgi(environ_overrides)
+        with flask.current_app.test_request_context(
+                "/login", environ_overrides=environ_overrides,
+            ):
+            # Now, we call Flask-WTF's method of generating a CSRF token...
+            csrf_token = generate_csrf()
+            # ...which also sets a value in `flask.session`, so we need to
+            # ask Flask to save that value to the cookie jar in the test
+            # client. This is where we actually use that request shim we made! 
+            flask.current_app.save_session(flask.session, request)
+            # And finally, return that CSRF token we got from Flask-WTF.
+            return csrf_token
+
+    # Feel free to define other methods on this test client. You can even
+    # use the `csrf_token` property we just defined, like we're doing here!
+    def login(self, email, password):
+        return self.post("/login", data={
+            "email": email,
+            "password": password,
+            "csrf_token": self.csrf_token,
+        }, follow_redirects=True)
+
+    def logout(self):
+        return self.get("/logout", follow_redirects=True)
+
+# To hook up this extended test client class to your Flask application,
+# assign it to the `test_client_class` property, like this:
+app = Flask(__name__)
+app.test_client_class = FlaskClient
+
+# Now in your tests, you can request a test client the same way
+# that you normally do:
+client = app.test_client()
+# But now, `client` is an instance of the class we defined!
+
+# In your tests, you can call the methods you defined, like this:
+client.login('user@example.com', 'passw0rd')
+
+# And any time you need to pass a CSRF token, just use the `csrf_token`
+# property, like this:
+client.post("/user/1", data={
+    "favorite_color": "blue",
+    "csrf_token": client.csrf_token,
+})