def set_cookie(
self,
server_name: str,
key: str,
value: str = "",
max_age: Optional[Union[int, timedelta]] = None,
expires: Optional[Union[int, float, datetime]] = None,
path: str = "/",
domain: Optional[str] = None,
secure: bool = False,
httponly: bool = False,
samesite: str = None,
charset: str = "utf-8",
) -> None:
"""Set a cookie in the cookie jar.
The arguments are the standard cookie morsels and this is a
wrapper around the stdlib SimpleCookie code.
"""
cookie = dump_cookie( # type: ignore
key,
value=value,
max_age=max_age,
expires=expires,
path=path,
domain=domain,
secure=secure,
httponly=httponly,
charset=charset,
samesite=samesite,
)
self.cookie_jar.extract_cookies(
_TestCookieJarResponse(Headers([("set-cookie", cookie)])), # type: ignore
U2Request(f"http://{server_name}{path}"),
)
def extract_wsgi(self, environ, headers):
"""Extract the server's set-cookie headers as cookies into the
cookie jar.
"""
self.extract_cookies(
_TestCookieResponse(headers), U2Request(get_current_url(environ))
)
async def session_transaction(
self,
path: str = "/",
*,
method: str = "GET",
headers: Optional[Union[dict, Headers]] = None,
query_string: Optional[dict] = None,
scheme: str = "http",
data: Optional[AnyStr] = None,
form: Optional[dict] = None,
json: Any = sentinel,
root_path: str = "",
http_version: str = "1.1",
) -> AsyncGenerator[Session, None]:
if self.cookie_jar is None:
raise RuntimeError("Session transactions only make sense with cookies enabled.")
headers, path, query_string_bytes = make_test_headers_path_and_query_string(
self.app, path, headers, query_string
)
request_body, body_headers = make_test_body_with_headers(data, form, json)
headers.update(**body_headers) # type: ignore
if self.cookie_jar is not None:
for cookie in self.cookie_jar:
headers.add("cookie", f"{cookie.name}={cookie.value}")
request = self.app.request_class(
method,
scheme,
path,
query_string_bytes,
headers,
root_path,
http_version,
send_push_promise=self._send_push_promise,
)
request.body.set_result(request_body)
original_request_ctx = _request_ctx_stack.top
async with self.app.request_context(request) as ctx: # type: ignore
session_interface = self.app.session_interface
session = await session_interface.open_session(self.app, ctx.request)
if session is None:
raise RuntimeError("Error opening the sesion. Check the secret_key?")
_request_ctx_stack.push(original_request_ctx)
try:
yield session
finally:
_request_ctx_stack.pop()
response = self.app.response_class(b"")
if not session_interface.is_null_session(session):
await session_interface.save_session(self.app, session, response)
self.cookie_jar.extract_cookies(
_TestCookieJarResponse(response.headers), # type: ignore
U2Request(ctx.request.url),
)
async def session_transaction(
self,
path: str = "/",
*,
method: str = "GET",
headers: Optional[Union[dict, Headers]] = None,
query_string: Optional[dict] = None,
scheme: str = "http",
data: Optional[AnyStr] = None,
form: Optional[dict] = None,
json: Any = sentinel,
root_path: str = "",
http_version: str = "1.1",
) -> AsyncGenerator[Session, None]:
if self.cookie_jar is None:
raise RuntimeError("Session transactions only make sense with cookies enabled.")
if headers is None:
headers = Headers()
elif isinstance(headers, Headers):
headers = headers
elif headers is not None:
headers = Headers(headers)
for cookie in self.cookie_jar:
headers.add("cookie", f"{cookie.name}={cookie.value}")
original_request_ctx = _request_ctx_stack.top
async with self.app.test_request_context(
path,
method=method,
headers=headers,
query_string=query_string,
scheme=scheme,
data=data,
form=form,
json=json,
root_path=root_path,
http_version=http_version,
) as ctx:
session_interface = self.app.session_interface
session = await session_interface.open_session(self.app, ctx.request)
if session is None:
raise RuntimeError("Error opening the sesion. Check the secret_key?")
_request_ctx_stack.push(original_request_ctx)
try:
yield session
finally:
_request_ctx_stack.pop()
response = self.app.response_class(b"")
if not session_interface.is_null_session(session):
await session_interface.save_session(self.app, session, response)
self.cookie_jar.extract_cookies(
_TestCookieJarResponse(response.headers), # type: ignore
U2Request(ctx.request.url),
)
async def _make_request(
self,
path: str,
method: str,
headers: Optional[Union[dict, Headers]],
data: Optional[AnyStr],
form: Optional[dict],
files: Optional[Dict[str, FileStorage]],
query_string: Optional[dict],
json: Any,
scheme: str,
root_path: str,
http_version: str,
scope_base: Optional[dict],
auth: Optional[Union[Authorization, Tuple[str, str]]] = None,
) -> Response:
headers, path, query_string_bytes = make_test_headers_path_and_query_string(
self.app, path, headers, query_string, auth)
request_data, body_headers = make_test_body_with_headers(data=data,
form=form,
files=files,
json=json,
app=self.app)
headers.update(**body_headers)
if self.cookie_jar is not None:
for cookie in self.cookie_jar:
headers.add("cookie", f"{cookie.name}={cookie.value}")
scope = make_test_scope(
"http",
path,
method,
headers,
query_string_bytes,
scheme,
root_path,
http_version,
scope_base,
_preserve_context=self.preserve_context,
)
async with self.http_connection_class(
self.app, scope,
_preserve_context=self.preserve_context) as connection:
await connection.send(request_data)
await connection.send_complete()
response = await connection.as_response()
if self.cookie_jar is not None:
self.cookie_jar.extract_cookies(
_TestCookieJarResponse(response.headers), # type: ignore
U2Request(f"{scheme}://{headers['host']}{path}"),
)
self.push_promises.extend(connection.push_promises)
return response
async def _make_request(
self,
path: str,
method: str = "GET",
headers: Optional[Union[dict, Headers]] = None,
data: Optional[AnyStr] = None,
form: Optional[dict] = None,
query_string: Optional[dict] = None,
json: Any = sentinel,
scheme: str = "http",
root_path: str = "",
http_version: str = "1.1",
) -> Response:
headers, path, query_string_bytes = make_test_headers_path_and_query_string(
self.app, path, headers, query_string)
request_data, body_headers = make_test_body_with_headers(
data, form, json, self.app)
# Replace with headers.update(**body_headers) when Werkzeug
# supports https://github.com/pallets/werkzeug/pull/1687
for key, value in body_headers.items():
headers[key] = value
if self.cookie_jar is not None:
for cookie in self.cookie_jar:
headers.add("cookie", f"{cookie.name}={cookie.value}")
request = self.app.request_class(
method,
scheme,
path,
query_string_bytes,
headers,
root_path,
http_version,
send_push_promise=self._send_push_promise,
)
request.body.set_result(request_data)
response = await self._handle_request(request)
if self.cookie_jar is not None:
self.cookie_jar.extract_cookies(
_TestCookieJarResponse(response.headers), # type: ignore
U2Request(request.url),
)
return response
def extract_asgi(self, scope, headers):
self.extract_cookies(
_TestCookieResponse(headers),
U2Request(get_current_url(scope)),
)
