From 8e454215dc27b2a31db49f6550d95d07a94ed0d8 Mon Sep 17 00:00:00 2001
From: c4ffein <c4ffein@gmail.com>
Date: Fri, 4 Oct 2024 15:38:09 +0200
Subject: [PATCH] Made NinjaTestClient accept COOKIES + documentation (#1311)

---
 docs/docs/guides/testing.md | 10 ++++++++++
 ninja/testing/client.py     |  7 +++++++
 tests/test_test_client.py   | 23 +++++++++++++++++++++++
 3 files changed, 40 insertions(+)

diff --git a/docs/docs/guides/testing.md b/docs/docs/guides/testing.md
index 9686cb05..4bfb31e9 100644
--- a/docs/docs/guides/testing.md
+++ b/docs/docs/guides/testing.md
@@ -38,6 +38,7 @@ It is also possible to access the deserialized data using the `data` property:
     self.assertEqual(response.data, {"msg": "Hello World"})
 ```
 
+## Attributes
 Arbitrary attributes can be added to the request object by passing keyword arguments to the client request methods:
 ```python
 class HelloTest(TestCase):
@@ -47,9 +48,18 @@ class HelloTest(TestCase):
         response = client.get("/hello", company_id=1)
 ```
 
+### Headers
 It is also possible to specify headers, both from the TestCase instanciation and the actual request:
 ```python
     client = TestClient(router, headers={"A": "a", "B": "b"})
     # The request will be made with {"A": "na", "B": "b", "C": "nc"} headers
     response = client.get("/test-headers", headers={"A": "na", "C": "nc"})
 ```
+
+### Cookies
+It is also possible to specify cookies, both from the TestCase instanciation and the actual request:
+```python
+    client = TestClient(router, COOKIES={"A": "a", "B": "b"})
+    # The request will be made with {"A": "na", "B": "b", "C": "nc"} cookies
+    response = client.get("/test-cookies", COOKIES={"A": "na", "C": "nc"})
+```
diff --git a/ninja/testing/client.py b/ninja/testing/client.py
index 8a3cbcd7..ecde0c23 100644
--- a/ninja/testing/client.py
+++ b/ninja/testing/client.py
@@ -30,8 +30,10 @@ def __init__(
         self,
         router_or_app: Union[NinjaAPI, Router],
         headers: Optional[Dict[str, str]] = None,
+        COOKIES: Optional[Dict[str, str]] = None,
     ) -> None:
         self.headers = headers or {}
+        self.cookies = COOKIES or {}
         self.router_or_app = router_or_app
 
     def get(
@@ -92,6 +94,11 @@ def request(
                 **self.headers,
                 **request_params.get("headers", {}),
             }
+        if self.cookies or request_params.get("COOKIES"):
+            request_params["COOKIES"] = {
+                **self.cookies,
+                **request_params.get("COOKIES", {}),
+            }
         func, request, kwargs = self._resolve(method, path, data, request_params)
         return self._call(func, request, kwargs)  # type: ignore
 
diff --git a/tests/test_test_client.py b/tests/test_test_client.py
index e182c20e..702cab75 100644
--- a/tests/test_test_client.py
+++ b/tests/test_test_client.py
@@ -32,6 +32,11 @@ def get_headers(request):
     return dict(request.headers)
 
 
+@router.get("/test-cookies")
+def get_cookies(request):
+    return dict(request.COOKIES)
+
+
 client = TestClient(router)
 
 
@@ -101,3 +106,21 @@ def test_headered_client_request_with_default_headers():
 def test_headered_client_request_with_overwritten_and_additional_headers():
     r = headered_client.get("/test-headers", headers={"A": "na", "C": "nc"})
     assert r.json() == {"A": "na", "B": "b", "C": "nc"}
+
+
+cookied_client = TestClient(router, COOKIES={"A": "a", "B": "b"})
+
+
+def test_client_request_only_cookies():
+    r = client.get("/test-cookies", COOKIES={"A": "na"})
+    assert r.json() == {"A": "na"}
+
+
+def test_headered_client_request_with_default_cookies():
+    r = cookied_client.get("/test-cookies")
+    assert r.json() == {"A": "a", "B": "b"}
+
+
+def test_headered_client_request_with_overwritten_and_additional_cookies():
+    r = cookied_client.get("/test-cookies", COOKIES={"A": "na", "C": "nc"})
+    assert r.json() == {"A": "na", "B": "b", "C": "nc"}