diff --git a/frontik/frontik_response.py b/frontik/frontik_response.py index e532be4dc..351891a2f 100644 --- a/frontik/frontik_response.py +++ b/frontik/frontik_response.py @@ -17,8 +17,13 @@ def __init__( body: bytes = b'', ): self.headers = HTTPHeaders(get_default_headers()) # type: ignore - if headers is not None: + + if isinstance(headers, HTTPHeaders): + for k, v in headers.get_all(): + self.headers.add(k, v) + elif headers is not None: self.headers.update(headers) + self.status_code = status_code self.body = body self.data_written = False diff --git a/tests/test_cookies.py b/tests/test_cookies.py new file mode 100644 index 000000000..689944bc7 --- /dev/null +++ b/tests/test_cookies.py @@ -0,0 +1,40 @@ +import pytest +from fastapi import Response + +from frontik.app import FrontikApplication +from frontik.handler import PageHandler, get_current_handler +from frontik.routing import plain_router +from frontik.testing import FrontikTestBase + + +@plain_router.get('/cookies', cls=PageHandler) +async def cookies_page(handler: PageHandler = get_current_handler()) -> None: + handler.set_cookie('key1', 'val1') + handler.set_cookie('key2', 'val2') + + +@plain_router.get('/asgi_cookies') +async def asgi_cookies_page(response: Response) -> None: + response.set_cookie('key1', 'val1') + response.set_cookie('key2', 'val2') + + +class TestFrontikTesting(FrontikTestBase): + @pytest.fixture(scope='class') + def frontik_app(self) -> FrontikApplication: + return FrontikApplication() + + async def test_cookies(self): + response = await self.fetch('/cookies') + + assert response.status_code == 200 + assert response.headers.getall('Set-Cookie') == ['key1=val1; Path=/', 'key2=val2; Path=/'] + + async def test_asgi_cookies(self): + response = await self.fetch('/asgi_cookies') + + assert response.status_code == 200 + assert response.headers.getall('Set-Cookie') == [ + 'key1=val1; Path=/; SameSite=lax', + 'key2=val2; Path=/; SameSite=lax', + ]