datasette/datasette/utils/testing.py

from asgiref.sync import async_to_sync
from urllib.parse import urlencode
import json

# These wrapper classes pre-date the introduction of
# datasette.client and httpx to Datasette. They could
# be removed if the Datasette tests are modified to
# call datasette.client directly.


class TestResponse:
    def __init__(self, httpx_response):
        self.httpx_response = httpx_response

    @property
    def status(self):
        return self.httpx_response.status_code

    # Supports both for test-writing convenience
    @property
    def status_code(self):
        return self.status

    @property
    def headers(self):
        return self.httpx_response.headers

    @property
    def body(self):
        return self.httpx_response.content

    @property
    def content(self):
        return self.body

    @property
    def cookies(self):
        return dict(self.httpx_response.cookies)

    @property
    def json(self):
        return json.loads(self.text)

    @property
    def text(self):
        return self.body.decode("utf8")


class TestClient:
    max_redirects = 5

    def __init__(self, ds):
        self.ds = ds

    def actor_cookie(self, actor):
        return self.ds.sign({"a": actor}, "actor")

    @async_to_sync
    async def get(
        self,
        path,
        follow_redirects=False,
        redirect_count=0,
        method="GET",
        cookies=None,
        if_none_match=None,
        headers=None,
    ):
        return await self._request(
            path=path,
            follow_redirects=follow_redirects,
            redirect_count=redirect_count,
            method=method,
            cookies=cookies,
            if_none_match=if_none_match,
            headers=headers,
        )

    @async_to_sync
    async def post(
        self,
        path,
        post_data=None,
        body=None,
        follow_redirects=False,
        redirect_count=0,
        content_type="application/x-www-form-urlencoded",
        cookies=None,
        headers=None,
        csrftoken_from=None,
    ):
        cookies = cookies or {}
        post_data = post_data or {}
        assert not (post_data and body), "Provide one or other of body= or post_data="
        # Maybe fetch a csrftoken first
        if csrftoken_from is not None:
            assert body is None, "body= is not compatible with csrftoken_from="
            if csrftoken_from is True:
                csrftoken_from = path
            token_response = await self._request(csrftoken_from, cookies=cookies)
            csrftoken = token_response.cookies["ds_csrftoken"]
            cookies["ds_csrftoken"] = csrftoken
            post_data["csrftoken"] = csrftoken
        if post_data:
            body = urlencode(post_data, doseq=True)
        return await self._request(
            path=path,
            follow_redirects=follow_redirects,
            redirect_count=redirect_count,
            method="POST",
            cookies=cookies,
            headers=headers,
            post_body=body,
            content_type=content_type,
        )

    @async_to_sync
    async def request(
        self,
        path,
        follow_redirects=True,
        redirect_count=0,
        method="GET",
        cookies=None,
        headers=None,
        post_body=None,
        content_type=None,
        if_none_match=None,
    ):
        return await self._request(
            path,
            follow_redirects=follow_redirects,
            redirect_count=redirect_count,
            method=method,
            cookies=cookies,
            headers=headers,
            post_body=post_body,
            content_type=content_type,
            if_none_match=if_none_match,
        )

    async def _request(
        self,
        path,
        follow_redirects=True,
        redirect_count=0,
        method="GET",
        cookies=None,
        headers=None,
        post_body=None,
        content_type=None,
        if_none_match=None,
    ):
        await self.ds.invoke_startup()
        headers = headers or {}
        if content_type:
            headers["content-type"] = content_type
        if if_none_match:
            headers["if-none-match"] = if_none_match
        httpx_response = await self.ds.client.request(
            method,
            path,
            follow_redirects=follow_redirects,
            avoid_path_rewrites=True,
            cookies=cookies,
            headers=headers,
            content=post_body,
        )
        response = TestResponse(httpx_response)
        if follow_redirects and response.status in (301, 302):
            assert (
                redirect_count < self.max_redirects
            ), f"Redirected {redirect_count} times, max_redirects={self.max_redirects}"
            location = response.headers["Location"]
            return await self._request(
                location, follow_redirects=True, redirect_count=redirect_count + 1
            )
        return response