Use csrftokens when running tests, refs #799

simonw · Jun 5, 2020 · 134296d · 134296d
1 parent 24c9377
commit 134296d
Show file tree

Hide file tree

Showing 3 changed files with 48 additions and 17 deletions.
diff --git a/tests/fixtures.py b/tests/fixtures.py
@@ -1,5 +1,5 @@
 from datasette.app import Datasette
-from datasette.utils import sqlite3
+from datasette.utils import sqlite3, MultiParams
 from asgiref.testing import ApplicationCommunicator
 from asgiref.sync import async_to_sync
 from http.cookies import SimpleCookie
@@ -60,10 +60,33 @@ async def get(
 
     @async_to_sync
     async def post(
-        self, path, post_data=None, allow_redirects=True, redirect_count=0, cookies=None
+        self,
+        path,
+        post_data=None,
+        allow_redirects=True,
+        redirect_count=0,
+        content_type="application/x-www-form-urlencoded",
+        cookies=None,
+        csrftoken_from=None,
     ):
+        cookies = cookies or {}
+        post_data = post_data or {}
+        # Maybe fetch a csrftoken first
+        if csrftoken_from is not None:
+            if csrftoken_from is True:
+                csrftoken_from = path
+            token_response = await self._request(csrftoken_from)
+            csrftoken = token_response.cookies["ds_csrftoken"]
+            cookies["ds_csrftoken"] = csrftoken
+            post_data["csrftoken"] = csrftoken
         return await self._request(
-            path, allow_redirects, redirect_count, "POST", cookies, post_data
+            path,
+            allow_redirects,
+            redirect_count,
+            "POST",
+            cookies,
+            post_data,
+            content_type,
         )
 
     async def _request(
@@ -74,6 +97,7 @@ async def _request(
         method="GET",
         cookies=None,
         post_data=None,
+        content_type=None,
     ):
         query_string = b""
         if "?" in path:
@@ -84,6 +108,8 @@ async def _request(
         else:
             raw_path = quote(path, safe="/:,").encode("latin-1")
         headers = [[b"host", b"localhost"]]
+        if content_type:
+            headers.append((b"content-type", content_type.encode("utf-8")))
         if cookies:
             sc = SimpleCookie()
             for key, value in cookies.items():
@@ -111,7 +137,7 @@ async def _request(
         start = await instance.receive_output(2)
         messages.append(start)
         assert start["type"] == "http.response.start"
-        headers = dict(
+        response_headers = MultiParams(
             [(k.decode("utf8"), v.decode("utf8")) for k, v in start["headers"]]
         )
         status = start["status"]
@@ -124,7 +150,7 @@ async def _request(
             body += message["body"]
             if not message.get("more_body"):
                 break
-        response = TestResponse(status, headers, body)
+        response = TestResponse(status, response_headers, body)
         if allow_redirects and response.status in (301, 302):
             assert (
                 redirect_count < self.max_redirects

diff --git a/tests/test_canned_write.py b/tests/test_canned_write.py
@@ -40,7 +40,7 @@ def canned_write_client():
 
 def test_insert(canned_write_client):
     response = canned_write_client.post(
-        "/data/add_name", {"name": "Hello"}, allow_redirects=False
+        "/data/add_name", {"name": "Hello"}, allow_redirects=False, csrftoken_from=True,
     )
     assert 302 == response.status
     assert "/data/add_name?success" == response.headers["Location"]
@@ -52,7 +52,7 @@ def test_insert(canned_write_client):
 
 def test_custom_success_message(canned_write_client):
     response = canned_write_client.post(
-        "/data/delete_name", {"rowid": 1}, allow_redirects=False
+        "/data/delete_name", {"rowid": 1}, allow_redirects=False, csrftoken_from=True
     )
     assert 302 == response.status
     messages = canned_write_client.ds.unsign(
@@ -62,11 +62,12 @@ def test_custom_success_message(canned_write_client):
 
 
 def test_insert_error(canned_write_client):
-    canned_write_client.post("/data/add_name", {"name": "Hello"})
+    canned_write_client.post("/data/add_name", {"name": "Hello"}, csrftoken_from=True)
     response = canned_write_client.post(
         "/data/add_name_specify_id",
         {"rowid": 1, "name": "Should fail"},
         allow_redirects=False,
+        csrftoken_from=True,
     )
     assert 302 == response.status
     assert "/data/add_name_specify_id?error" == response.headers["Location"]
@@ -82,6 +83,7 @@ def test_insert_error(canned_write_client):
         "/data/add_name_specify_id",
         {"rowid": 1, "name": "Should fail"},
         allow_redirects=False,
+        csrftoken_from=True,
     )
     assert [["ERROR", 3]] == canned_write_client.ds.unsign(
         response.cookies["ds_messages"], "messages"

diff --git a/tests/test_utils.py b/tests/test_utils.py
@@ -439,15 +439,18 @@ def foo(a, b):
         utils.call_with_supported_arguments(foo, a=1)
 
 
-@pytest.mark.parametrize("data,should_raise", [
-    ([["foo", "bar"], ["foo", "baz"]], False),
-    ([("foo", "bar"), ("foo", "baz")], False),
-    ((["foo", "bar"], ["foo", "baz"]), False),
-    ([["foo", "bar"], ["foo", "baz", "bax"]], True),
-    ({"foo": ["bar", "baz"]}, False),
-    ({"foo": ("bar", "baz")}, False),
-    ({"foo": "bar"}, True),
-])
+@pytest.mark.parametrize(
+    "data,should_raise",
+    [
+        ([["foo", "bar"], ["foo", "baz"]], False),
+        ([("foo", "bar"), ("foo", "baz")], False),
+        ((["foo", "bar"], ["foo", "baz"]), False),
+        ([["foo", "bar"], ["foo", "baz", "bax"]], True),
+        ({"foo": ["bar", "baz"]}, False),
+        ({"foo": ("bar", "baz")}, False),
+        ({"foo": "bar"}, True),
+    ],
+)
 def test_multi_params(data, should_raise):
     if should_raise:
         with pytest.raises(AssertionError):