Add SetCookie utility (#254)

Author: lundberg

docs/api.md

@@ -133,15 +133,17 @@ Setter for the [side effect](guide.md#mock-with-a-side-effect) to trigger.
 
 Shortcut for creating and mocking a `HTTPX` [Response](#response).
 
-> <code>route.<strong>respond</strong>(*status_code=200, headers=None, content=None, text=None, html=None, json=None, stream=None*)</strong></code>
+> <code>route.<strong>respond</strong>(*status_code=200, headers=None, cookies=None, content=None, text=None, html=None, json=None, stream=None, content_type=None*)</strong></code>
 >
 > **Parameters:**
 >
 > * **status_code** - *(optional) int - default: `200`*  
 >   Response status code to mock.
-> * **headers** - *(optional) dict*  
+> * **headers** - *(optional) dict | Sequence[tuple[str, str]]*  
 >   Response headers to mock.
-> * **content** - *(optional) bytes | str | iterable bytes*  
+> * **cookies** - *(optional) dict | Sequence[tuple[str, str]] | Sequence[SetCookie]*  
+>   Response cookies to mock as `Set-Cookie` headers. See [SetCookie](#setcookie).
+> * **content** - *(optional) bytes | str | Iterable[bytes]*  
 >   Response raw content to mock.
 > * **text** - *(optional) str*  
 >   Response *text* content to mock, with automatic content-type header added.
@@ -151,6 +153,8 @@ Shortcut for creating and mocking a `HTTPX` [Response](#response).
 >   Response *JSON* content to mock, with automatic content-type header added.
 > * **stream** - *(optional) Iterable[bytes]*  
 >   Response *stream* to mock.
+> * **content_type** - *(optional) str*  
+>   Response `Content-Type` header to mock.
 >
 > **Returns:** `Route`
 
@@ -191,6 +195,24 @@ Shortcut for creating and mocking a `HTTPX` [Response](#response).
 > * **stream** - *(optional) Iterable[bytes]*  
 >   Content *stream*.
 
+!!! tip "Cookies"
+    Use [respx.SetCookie(...)](#setcookie) to produce `Set-Cookie` headers.
+
+---
+
+## SetCookie
+
+A utility to render a `("Set-Cookie", <cookie header value>)` tuple. See route [respond](#respond) shortcut for alternative use.
+
+> <code>respx.<strong>SetCookie</strong>(*name, value, path=None, domain=None, expires=None, max_age=None, http_only=False, same_site=None, secure=False, partitioned=False*)</strong></code>
+
+``` python
+import respx
+respx.post("https://example.org/").mock(
+    return_value=httpx.Response(200, headers=[SetCookie("foo", "bar")])
+)
+```
+
 ---
 
 ## Patterns

respx/__init__.py

@@ -2,6 +2,7 @@
 from .handlers import ASGIHandler, WSGIHandler
 from .models import MockResponse, Route
 from .router import MockRouter, Router
+from .utils import SetCookie
 
 from .api import (  # isort:skip
     mock,
@@ -24,6 +25,7 @@
     options,
 )
 
+
 __all__ = [
     "__version__",
     "MockResponse",
@@ -32,6 +34,7 @@
     "WSGIHandler",
     "Router",
     "Route",
+    "SetCookie",
     "mock",
     "routes",
     "calls",

respx/models.py

@@ -16,10 +16,13 @@
 
 import httpx
 
+from respx.utils import SetCookie
+
 from .patterns import M, Pattern
 from .types import (
     CallableSideEffect,
     Content,
+    CookieTypes,
     HeaderTypes,
     ResolvedResponseTypes,
     RouteResultTypes,
@@ -90,6 +93,7 @@ def __init__(
         content: Optional[Content] = None,
         content_type: Optional[str] = None,
         http_version: Optional[str] = None,
+        cookies: Optional[Union[CookieTypes, Sequence[SetCookie]]] = None,
         **kwargs: Any,
     ) -> None:
         if not isinstance(content, (str, bytes)) and (
@@ -110,6 +114,19 @@ def __init__(
         if content_type:
             self.headers["Content-Type"] = content_type
 
+        if cookies:
+            if isinstance(cookies, dict):
+                cookies = tuple(cookies.items())
+            self.headers = httpx.Headers(
+                (
+                    *self.headers.multi_items(),
+                    *(
+                        cookie if isinstance(cookie, SetCookie) else SetCookie(*cookie)
+                        for cookie in cookies
+                    ),
+                )
+            )
+
 
 class Route:
     def __init__(
@@ -256,6 +273,7 @@ def respond(
         status_code: int = 200,
         *,
         headers: Optional[HeaderTypes] = None,
+        cookies: Optional[Union[CookieTypes, Sequence[SetCookie]]] = None,
         content: Optional[Content] = None,
         text: Optional[str] = None,
         html: Optional[str] = None,
@@ -268,6 +286,7 @@ def respond(
         response = MockResponse(
             status_code,
             headers=headers,
+            cookies=cookies,
             content=content,
             text=text,
             html=html,

respx/utils.py

@@ -1,8 +1,14 @@
 import email
+from datetime import datetime
 from email.message import Message
-from typing import List, Tuple, cast
+from typing import Dict, List, NamedTuple, Optional, Tuple, Type, TypeVar, Union, cast
 from urllib.parse import parse_qsl
 
+try:
+    from typing import Literal  # type: ignore[attr-defined]
+except ImportError:  # pragma: no cover
+    from typing_extensions import Literal
+
 import httpx
 
 
@@ -71,3 +77,62 @@ def decode_data(request: httpx.Request) -> Tuple[MultiItems, MultiItems]:
         files = MultiItems()
 
     return data, files
+
+
+Self = TypeVar("Self", bound="SetCookie")
+
+
+class SetCookie(
+    NamedTuple(
+        "SetCookie",
+        [
+            ("header_name", Literal["Set-Cookie"]),
+            ("header_value", str),
+        ],
+    )
+):
+    def __new__(
+        cls: Type[Self],
+        name: str,
+        value: str,
+        *,
+        path: Optional[str] = None,
+        domain: Optional[str] = None,
+        expires: Optional[Union[str, datetime]] = None,
+        max_age: Optional[int] = None,
+        http_only: bool = False,
+        same_site: Optional[Literal["Strict", "Lax", "None"]] = None,
+        secure: bool = False,
+        partitioned: bool = False,
+    ) -> Self:
+        """
+        https://developer.mozilla.org/en-US/docs/Web/HTTP/Headers/Set-Cookie#syntax
+        """
+        attrs: Dict[str, Union[str, bool]] = {name: value}
+        if path is not None:
+            attrs["Path"] = path
+        if domain is not None:
+            attrs["Domain"] = domain
+        if expires is not None:
+            if isinstance(expires, datetime):  # pragma: no branch
+                expires = expires.strftime("%a, %d %b %Y %H:%M:%S GMT")
+            attrs["Expires"] = expires
+        if max_age is not None:
+            attrs["Max-Age"] = str(max_age)
+        if http_only:
+            attrs["HttpOnly"] = True
+        if same_site is not None:
+            attrs["SameSite"] = same_site
+            if same_site == "None":  # pragma: no branch
+                secure = True
+        if secure:
+            attrs["Secure"] = True
+        if partitioned:
+            attrs["Partitioned"] = True
+
+        string = "; ".join(
+            _name if _value is True else f"{_name}={_value}"
+            for _name, _value in attrs.items()
+        )
+        self = super().__new__(cls, "Set-Cookie", string)
+        return self

tests/test_api.py

@@ -564,6 +564,46 @@ def test_respond():
             route.respond(content=Exception())  # type: ignore[arg-type]
 
 
+def test_can_respond_with_cookies():
+    with respx.mock:
+        route = respx.get("https://foo.bar/").respond(
+            json={}, headers={"X-Foo": "bar"}, cookies={"foo": "bar", "ham": "spam"}
+        )
+        response = httpx.get("https://foo.bar/")
+        assert len(response.headers) == 5
+        assert response.headers["X-Foo"] == "bar", "mocked header is missing"
+        assert len(response.cookies) == 2
+        assert response.cookies["foo"] == "bar"
+        assert response.cookies["ham"] == "spam"
+
+        route.respond(cookies=[("egg", "yolk")])
+        response = httpx.get("https://foo.bar/")
+        assert len(response.cookies) == 1
+        assert response.cookies["egg"] == "yolk"
+
+        route.respond(
+            cookies=[respx.SetCookie("foo", "bar", path="/", same_site="Lax")]
+        )
+        response = httpx.get("https://foo.bar/")
+        assert len(response.cookies) == 1
+        assert response.cookies["foo"] == "bar"
+
+
+def test_can_mock_response_with_set_cookie_headers():
+    request = httpx.Request("GET", "https://example.com/")
+    response = httpx.Response(
+        200,
+        headers=[
+            respx.SetCookie("foo", value="bar"),
+            respx.SetCookie("ham", value="spam"),
+        ],
+        request=request,
+    )
+    assert len(response.cookies) == 2
+    assert response.cookies["foo"] == "bar"
+    assert response.cookies["ham"] == "spam"
+
+
 @pytest.mark.parametrize(
     "kwargs",
     [

tests/test_utils.py

@@ -0,0 +1,33 @@
+from datetime import datetime, timezone
+
+from respx.utils import SetCookie
+
+
+class TestSetCookie:
+    def test_can_render_all_attributes(self) -> None:
+        expires = datetime.fromtimestamp(0, tz=timezone.utc)
+        cookie = SetCookie(
+            "foo",
+            value="bar",
+            path="/",
+            domain=".example.com",
+            expires=expires,
+            max_age=44,
+            http_only=True,
+            same_site="None",
+            partitioned=True,
+        )
+        assert cookie == (
+            "Set-Cookie",
+            (
+                "foo=bar; "
+                "Path=/; "
+                "Domain=.example.com; "
+                "Expires=Thu, 01 Jan 1970 00:00:00 GMT; "
+                "Max-Age=44; "
+                "HttpOnly; "
+                "SameSite=None; "
+                "Secure; "
+                "Partitioned"
+            ),
+        )