Test Client¶

def build_request(
    self,
    method: str,
    url: URL | str,
    *,
    content: RequestContent | None = None,
    data: RequestData | None = None,
    files: RequestFiles | None = None,
    json: typing.Any | None = None,
    params: QueryParamTypes | None = None,
    headers: HeaderTypes | None = None,
    cookies: CookieTypes | None = None,
    timeout: TimeoutTypes | UseClientDefault = USE_CLIENT_DEFAULT,
    extensions: RequestExtensions | None = None,
) -> Request:
    """
    Build and return a request instance.

    * The `params`, `headers` and `cookies` arguments
    are merged with any values set on the client.
    * The `url` argument is merged with any `base_url` set on the client.

    See also: [Request instances][0]

    [0]: /advanced/clients/#request-instances
    """
    url = self._merge_url(url)
    headers = self._merge_headers(headers)
    cookies = self._merge_cookies(cookies)
    params = self._merge_queryparams(params)
    extensions = {} if extensions is None else extensions
    if "timeout" not in extensions:
        timeout = (
            self.timeout
            if isinstance(timeout, UseClientDefault)
            else Timeout(timeout)
        )
        extensions = dict(**extensions, timeout=timeout.as_dict())
    return Request(
        method,
        url,
        content=content,
        data=data,
        files=files,
        json=json,
        params=params,
        headers=headers,
        cookies=cookies,
        extensions=extensions,
    )

def request(
    self,
    method: str,
    url: URLTypes,
    *,
    content: RequestContent | None = None,
    data: RequestData | None = None,
    files: RequestFiles | None = None,
    json: Any = None,
    params: QueryParamTypes | None = None,
    headers: HeaderTypes | None = None,
    cookies: CookieTypes | None = None,
    auth: (AuthTypes | httpx._client.UseClientDefault) = httpx._client.USE_CLIENT_DEFAULT,
    follow_redirects: bool | None = None,
    timeout: (
        TimeoutTypes | httpx._client.UseClientDefault
    ) = httpx._client.USE_CLIENT_DEFAULT,
    extensions: dict[str, Any] | None = None,
    stream: bool = False,
) -> httpx.Response:
    """
    Builds and sends a network request to the ASGI application.

    This method overrides the standard `httpx.Client.request` to provide specific
    handling for the TestClient, including URL merging and stream handling.

    Args:
        method (str): The HTTP method to use (e.g., 'GET', 'POST').
        url (URLTypes): The URL to send the request to.
        content (RequestContent | None, optional): Binary content to send in the body.
            Defaults to None.
        data (RequestData | None, optional): Form data to send in the body.
            Defaults to None.
        files (RequestFiles | None, optional): Files to upload. Defaults to None.
        json (Any, optional): JSON data to send in the body. Defaults to None.
        params (QueryParamTypes | None, optional): Query parameters to include in the URL.
            Defaults to None.
        headers (HeaderTypes | None, optional): Headers to include in the request.
            Defaults to None.
        cookies (CookieTypes | None, optional): Cookies to include in the request.
            Defaults to None.
        auth (AuthTypes | httpx._client.UseClientDefault, optional): Authentication
            credentials. Defaults to httpx._client.USE_CLIENT_DEFAULT.
        follow_redirects (bool | None, optional): Whether to follow redirects.
            Defaults to None.
        timeout (TimeoutTypes | httpx._client.UseClientDefault, optional): Timeout
            configuration. Defaults to httpx._client.USE_CLIENT_DEFAULT.
        extensions (dict[str, Any] | None, optional): Extensions to include in the request.
            Defaults to None.
        stream (bool, optional): Whether to stream the response content. Defaults to False.

    Returns:
        httpx.Response: The response received from the application.
    """
    url = self._merge_url(url)
    redirect: bool | httpx._client.UseClientDefault = httpx._client.USE_CLIENT_DEFAULT
    if follow_redirects is not None:
        redirect = follow_redirects

    # httpx 0.28+ deprecated per-request cookies= parameter
    # Save current cookies, merge per-request cookies, then restore after request
    if cookies is not None:
        _saved_cookies = dict(self.cookies)
        self.cookies.update(cookies)
    else:
        _saved_cookies = None

    try:
        if stream:
            return self.stream(
                method=method,
                url=url,
                content=content,
                data=data,
                files=files,
                json=json,
                params=params,
                headers=headers,
                auth=auth,
                timeout=timeout,
                extensions=extensions,
                follow_redirects=follow_redirects,
            ).__enter__()

        return super().request(
            method,
            url,
            content=content,
            data=data,
            files=files,
            json=json,
            params=params,
            headers=headers,
            auth=auth,
            follow_redirects=redirect,
            timeout=timeout,
            extensions=extensions,
        )
    finally:
        if _saved_cookies is not None:
            self.cookies.clear()
            self.cookies.update(_saved_cookies)

@contextmanager
def stream(
    self,
    method: str,
    url: URL | str,
    *,
    content: RequestContent | None = None,
    data: RequestData | None = None,
    files: RequestFiles | None = None,
    json: typing.Any | None = None,
    params: QueryParamTypes | None = None,
    headers: HeaderTypes | None = None,
    cookies: CookieTypes | None = None,
    auth: AuthTypes | UseClientDefault | None = USE_CLIENT_DEFAULT,
    follow_redirects: bool | UseClientDefault = USE_CLIENT_DEFAULT,
    timeout: TimeoutTypes | UseClientDefault = USE_CLIENT_DEFAULT,
    extensions: RequestExtensions | None = None,
) -> typing.Iterator[Response]:
    """
    Alternative to `httpx.request()` that streams the response body
    instead of loading it into memory at once.

    **Parameters**: See `httpx.request`.

    See also: [Streaming Responses][0]

    [0]: /quickstart#streaming-responses
    """
    request = self.build_request(
        method=method,
        url=url,
        content=content,
        data=data,
        files=files,
        json=json,
        params=params,
        headers=headers,
        cookies=cookies,
        timeout=timeout,
        extensions=extensions,
    )
    response = self.send(
        request=request,
        auth=auth,
        follow_redirects=follow_redirects,
        stream=True,
    )
    try:
        yield response
    finally:
        response.close()

def send(
    self,
    request: Request,
    *,
    stream: bool = False,
    auth: AuthTypes | UseClientDefault | None = USE_CLIENT_DEFAULT,
    follow_redirects: bool | UseClientDefault = USE_CLIENT_DEFAULT,
) -> Response:
    """
    Send a request.

    The request is sent as-is, unmodified.

    Typically you'll want to build one with `Client.build_request()`
    so that any client-level configuration is merged into the request,
    but passing an explicit `httpx.Request()` is supported as well.

    See also: [Request instances][0]

    [0]: /advanced/clients/#request-instances
    """
    if self._state == ClientState.CLOSED:
        raise RuntimeError("Cannot send a request, as the client has been closed.")

    self._state = ClientState.OPENED
    follow_redirects = (
        self.follow_redirects
        if isinstance(follow_redirects, UseClientDefault)
        else follow_redirects
    )

    self._set_timeout(request)

    auth = self._build_request_auth(request, auth)

    response = self._send_handling_auth(
        request,
        auth=auth,
        follow_redirects=follow_redirects,
        history=[],
    )
    try:
        if not stream:
            response.read()

        return response

    except BaseException as exc:
        response.close()
        raise exc

def get(self, url: URLTypes, **kwargs: Any) -> httpx.Response:
    """
    Sends a GET request.

    Args:
        url (URLTypes): The URL to send the request to.
        **kwargs (Any): Additional arguments passed to `request`.

    Returns:
        httpx.Response: The response from the server.
    """
    return self._process_request(method="GET", url=url, **kwargs)  # type: ignore

def options(self, url: URLTypes, **kwargs: Any) -> httpx.Response:
    """
    Sends an OPTIONS request.

    Args:
        url (URLTypes): The URL to send the request to.
        **kwargs (Any): Additional arguments passed to `request`.

    Returns:
        httpx.Response: The response from the server.
    """
    return self._process_request(method="OPTIONS", url=url, **kwargs)  # type: ignore

def head(self, url: URLTypes, **kwargs: Any) -> httpx.Response:
    """
    Sends a HEAD request.

    Args:
        url (URLTypes): The URL to send the request to.
        **kwargs (Any): Additional arguments passed to `request`.

    Returns:
        httpx.Response: The response from the server.
    """
    return self._process_request(method="HEAD", url=url, **kwargs)  # type: ignore

def post(self, url: URLTypes, **kwargs: Any) -> httpx.Response:
    """
    Sends a POST request.

    Args:
        url (URLTypes): The URL to send the request to.
        **kwargs (Any): Additional arguments passed to `request`.

    Returns:
        httpx.Response: The response from the server.
    """
    return self._process_request(method="POST", url=url, **kwargs)  # type: ignore

def put(self, url: URLTypes, **kwargs: Any) -> httpx.Response:
    """
    Sends a PUT request.

    Args:
        url (URLTypes): The URL to send the request to.
        **kwargs (Any): Additional arguments passed to `request`.

    Returns:
        httpx.Response: The response from the server.
    """
    return self._process_request(method="PUT", url=url, **kwargs)  # type: ignore

def patch(self, url: URLTypes, **kwargs: Any) -> httpx.Response:
    """
    Sends a PATCH request.

    Args:
        url (URLTypes): The URL to send the request to.
        **kwargs (Any): Additional arguments passed to `request`.

    Returns:
        httpx.Response: The response from the server.
    """
    return self._process_request(method="PATCH", url=url, **kwargs)  # type: ignore

def delete(self, url: URLTypes, **kwargs: Any) -> httpx.Response:
    """
    Sends a DELETE request.

    Args:
        url (URLTypes): The URL to send the request to.
        **kwargs (Any): Additional arguments passed to `request`.

    Returns:
        httpx.Response: The response from the server.
    """
    return self._process_request(method="DELETE", url=url, **kwargs)  # type: ignore

def close(self) -> None:
    """
    Close transport and proxies.
    """
    if self._state != ClientState.CLOSED:
        self._state = ClientState.CLOSED

        self._transport.close()
        for transport in self._mounts.values():
            if transport is not None:
                transport.close()

def authenticate(self, user: Any) -> TestClient:
    """
    Sets an authenticated user in the application state.

    This helper method injects a user object into the application's state, simulating
    a logged-in user session for subsequent requests.

    Args:
        user (Any): The user object to be set as authenticated.

    Returns:
        TestClient: The instance of the client, allowing for method chaining.
    """
    self.app_state[_AUTH_USER_KEY] = user
    return self

def logout(self) -> TestClient:
    """
    Removes the authenticated user from the application state.

    This helper method clears any user object currently stored in the application's
    state, simulating a logout action.

    Returns:
        TestClient: The instance of the client, allowing for method chaining.
    """
    self.app_state.pop(_AUTH_USER_KEY, None)
    return self

@contextlib.contextmanager
def authenticated(self, user: Any) -> Generator[TestClient, None, None]:
    """
    A context manager that temporarily authenticates a user for the duration of the block.

    This is useful for testing protected endpoints without permanently altering the
    client's authentication state. The previous authentication state is restored
    when the context exits.

    Args:
        user (Any): The user object to be temporarily set as authenticated.

    Yields:
        TestClient: The authenticated client instance.

    Example:
        async with AsyncTestClient(app) as client:
            with client.authenticated(user):
                response = await client.get("/protected")
    """
    previous = self.app_state.get(_AUTH_USER_KEY)
    self.app_state[_AUTH_USER_KEY] = user
    try:
        yield self
    finally:
        if previous is None:
            self.app_state.pop(_AUTH_USER_KEY, None)
        else:
            self.app_state[_AUTH_USER_KEY] = previous

def websocket_connect(
    self,
    url: str,
    subprotocols: Sequence[str] | None = None,
    **kwargs: Any,
) -> WebSocketTestSession:
    """
    Initiates a WebSocket connection to the application.

    This method constructs the necessary headers for a WebSocket upgrade request and
    establishes a session if the upgrade is successful.

    Args:
        url (str): The URL path for the WebSocket connection (relative to base_url).
        subprotocols (Sequence[str] | None, optional): A list of WebSocket subprotocols
            to request. Defaults to None.
        **kwargs (Any): Additional arguments passed to the request.

    Returns:
        WebSocketTestSession: An active WebSocket session object.

    Raises:
        RuntimeError: If the server does not upgrade the connection to a WebSocket.
    """
    url = urljoin("ws://testserver", url)
    headers = self._prepare_websocket_headers(subprotocols, **kwargs)
    kwargs["headers"] = headers
    try:
        super().request("GET", url, **kwargs)
    except UpgradeException as exc:
        session = exc.session
    else:
        raise RuntimeError("Expected WebSocket upgrade")

    return session

async def lifespan(self) -> None:
    """
    Manages the ASGI lifespan protocol loop.

    This method sends the lifespan context to the app and handles the receive/send
    channels for lifespan events.
    """
    scope = {"type": "lifespan", "state": self.app_state}
    try:
        await self.app(scope, self.stream_receive.receive, self.stream_send.send)
    finally:
        with contextlib.suppress(anyio.ClosedResourceError):
            await self.stream_send.send(None)

async def wait_startup(self) -> None:
    """
    Waits for the ASGI application to complete its startup sequence.

    Sends the 'lifespan.startup' event and awaits the confirmation message.

    Raises:
        AssertionError: If the received message type is not a valid startup response.
    """
    await self.stream_receive.send({"type": "lifespan.startup"})

    async def receive() -> Any:
        message = await self.stream_send.receive()
        if message is None:
            self.task.result()
        return message

    message = await receive()
    assert message["type"] in (
        "lifespan.startup.complete",
        "lifespan.startup.failed",
    )
    if message["type"] == "lifespan.startup.failed":
        await receive()

async def wait_shutdown(self) -> None:
    """
    Waits for the ASGI application to complete its shutdown sequence.

    Sends the 'lifespan.shutdown' event and awaits the confirmation message.
    Handles both sync (portal) and async (task group) modes.

    Raises:
        AssertionError: If the received message type is not a valid shutdown response.
    """

    async def receive() -> Any:
        message = await self.stream_send.receive()
        if message is None:
            self.task.result()
        return message

    async_mode = hasattr(self, "_tg")

    if async_mode:
        await self.stream_receive.send({"type": "lifespan.shutdown"})
        message = await receive()
        assert message["type"] in (
            "lifespan.shutdown.complete",
            "lifespan.shutdown.failed",
        )
        if message["type"] == "lifespan.shutdown.failed":
            await receive()
    else:
        async with self.stream_send:
            await self.stream_receive.send({"type": "lifespan.shutdown"})
            message = await receive()
            assert message["type"] in (
                "lifespan.shutdown.complete",
                "lifespan.shutdown.failed",
            )
            if message["type"] == "lifespan.shutdown.failed":
                await receive()