✨ ASGI Cookies component

vortico · Sep 3, 2024 · 0695c77 · 0695c77
1 parent 4f47f49
commit 0695c77
Show file tree

Hide file tree

Showing 3 changed files with 115 additions and 3 deletions.
diff --git a/flama/asgi.py b/flama/asgi.py
@@ -1,6 +1,7 @@
+from http.cookies import SimpleCookie
 from urllib.parse import parse_qsl
 
-from flama import types
+from flama import http, types
 from flama.injection.components import Component, Components
 
 __all__ = [
@@ -66,8 +67,17 @@ def resolve(self, scope: types.Scope) -> types.QueryParams:
 
 
 class HeadersComponent(Component):
-    def resolve(self, scope: types.Scope) -> types.Headers:
-        return types.Headers(scope=scope)
+    def resolve(self, request: http.Request) -> types.Headers:
+        return request.headers
+
+
+class CookiesComponent(Component):
+    def resolve(self, headers: types.Headers) -> types.Cookies:
+        cookie = SimpleCookie()
+        cookie.load(headers.get("cookie", ""))
+        return types.Cookies(
+            {str(name): {str(k): str(v) for k, v in morsel.items()} for name, morsel in cookie.items()}
+        )
 
 
 class BodyComponent(Component):
@@ -95,6 +105,7 @@ async def resolve(self, receive: types.Receive) -> types.Body:
         QueryStringComponent(),
         QueryParamsComponent(),
         HeadersComponent(),
+        CookiesComponent(),
         BodyComponent(),
     ]
 )
diff --git a/flama/types/http.py b/flama/types/http.py
@@ -23,6 +23,7 @@
     "URL",
     "Headers",
     "MutableHeaders",
+    "Cookies",
     "QueryParams",
     "PARAMETERS_TYPES",
 ]
@@ -41,6 +42,7 @@
 RequestData = t.NewType("RequestData", t.Dict[str, t.Any])
 Headers = starlette.datastructures.Headers
 MutableHeaders = starlette.datastructures.MutableHeaders
+Cookies = t.NewType("Cookies", t.Dict[str, t.Dict[str, str]])
 QueryParams = starlette.datastructures.QueryParams
 
 PARAMETERS_TYPES: t.Dict[t.Type, t.Type] = {

diff --git a/tests/test_asgi.py b/tests/test_asgi.py
@@ -1,3 +1,5 @@
+import http.cookiejar
+
 import pytest
 
 from flama import types
@@ -301,6 +303,103 @@ async def test_headers(self, client, path, method, request_kwargs, expected):
         assert response_json == expected
 
 
+class TestCaseCookiesComponent:
+    @pytest.fixture(scope="function", autouse=True)
+    def add_endpoints(self, app):
+        @app.route("/cookies/", methods=["GET", "POST"])
+        def get_cookies(cookies: types.Cookies):
+            return {"cookies": dict(cookies)}
+
+    @pytest.mark.parametrize(
+        ["path", "method", "cookies", "expected"],
+        [
+            pytest.param(
+                "http://example.com/cookies/",
+                "get",
+                {},
+                {"cookies": {}},
+                id="default",
+            ),
+            pytest.param(
+                "http://example.com/cookies/",
+                "get",
+                [
+                    http.cookiejar.Cookie(
+                        version=0,
+                        name="foo",
+                        value="bar",
+                        port=None,
+                        port_specified=False,
+                        domain="",
+                        domain_specified=False,
+                        domain_initial_dot=False,
+                        path="/",
+                        path_specified=True,
+                        secure=False,
+                        expires=None,
+                        discard=True,
+                        comment=None,
+                        comment_url=None,
+                        rest={"HttpOnly": ""},
+                        rfc2109=False,
+                    )
+                ],
+                {
+                    "cookies": {
+                        "foo": {
+                            "expires": "",
+                            "path": "",
+                            "comment": "",
+                            "domain": "",
+                            "max-age": "",
+                            "secure": "",
+                            "httponly": "",
+                            "version": "",
+                            "samesite": "",
+                        }
+                    }
+                },
+                id="cookie",
+            ),
+            pytest.param(
+                "http://example.com/cookies/",
+                "get",
+                [
+                    http.cookiejar.Cookie(
+                        version=0,
+                        name="foo",
+                        value="bar",
+                        port=None,
+                        port_specified=False,
+                        domain="",
+                        domain_specified=False,
+                        domain_initial_dot=False,
+                        path="/",
+                        path_specified=True,
+                        secure=True,
+                        expires=None,
+                        discard=True,
+                        comment=None,
+                        comment_url=None,
+                        rest={"HttpOnly": "true"},
+                        rfc2109=False,
+                    )
+                ],
+                {"cookies": {}},  # Cannot get cookie because secure and no https
+                id="cookie_secure",
+            ),
+        ],
+    )
+    async def test_cookies(self, client, path, method, cookies, expected):
+        cookies_jar = http.cookiejar.CookieJar()
+        for cookie in cookies:
+            cookies_jar.set_cookie(cookie)
+        response = await client.request(method, path, cookies=cookies_jar)
+        response_json = response.json()
+
+        assert response_json == expected
+
+
 class TestCaseBodyComponent:
     @pytest.fixture(scope="function", autouse=True)
     def add_endpoints(self, app):