Testing with auth middleware and guards

[aranvir]

Hi, I had a first look at testing with litestar and how to set it up for an app with an auth middleware. My rationale is that I want to be able to test the endpoints functionally without considering authentication (where possible) and then have a separate set of tests to confirm proper authentication and authorization behavior.

The problem

Switching between auth-enabled and auth-disabled testing was simple by using either an "app client" (AsyncTestClient(app=app)) or a "controller client" (create_async_test_client(UserController)). Figuring out the right configuration for testing an "admin user guard" was not as straightforward.

I based my work on the docs for guards and testing using sessions. Thus, I expected something like this to work

...
    async with create_async_test_client(
        UserController,
        session_config=session_config,
        debug=True
    ) as client:
        await client.set_session_data({"user_name": "admin"})
        yield client
...

However this solution complains that the auth middleware is missing (fair, the session middleware does not specify a user retrieve handler). Unexpectedly, just setting the on_app_init parameter as one does for setting up the app also does not work - it's missing the session configuration (which the session_auth object owns).

...
    async with create_async_test_client(
        UserController,
        on_app_init=[session_auth.on_app_init],
        debug=True
    ) as client:
        await client.set_session_data({"user_name": "admin"})
        yield client
...

What ended up working was the following: Both the session_auth and the session_config object need to be provided for creating a test client which is in contrast to how the app is set up (where the session_config object is passed to the session_auth object and that then installs the session middleware).

...
    async with create_async_test_client(
        UserController,
        session_config=session_config,
        on_app_init=[session_auth.on_app_init],
        debug=True
    ) as client:
        await client.set_session_data({"user_name": "admin"})
        yield client
...

My questions

• Is there a good reason why the test client needs to receive session_auth and session_config separately?
• Is this the right approach to test endpoints with and without authentication. Because technically the user_client still uses the same authentication and authorization mechanisms as the app - it just bypasses the login and is always set to an admin user.

The Code

Full app and test code below. I was running this on Win10, Python 3.11.7, with Litestar 2.8.2.
App

from typing import TYPE_CHECKING, Any
from dataclasses import dataclass

from litestar import Litestar, Controller, get, post, exceptions, Request
from litestar.security.session_auth.auth import SessionAuth
from litestar.middleware.session.server_side import ServerSideSessionConfig, ServerSideSessionBackend

if TYPE_CHECKING:
    from litestar.connection import ASGIConnection
    from litestar.handlers.base import BaseRouteHandler


@dataclass
class User:
    name: str
    pw: str
    admin: bool


MOCK_DB = {
    "dummy": User("dummy", "test", False),
    "admin": User("admin", "admin", True)
}


def admin_user_guard(connection: "ASGIConnection", _: "BaseRouteHandler") -> None:
    """Access is only permitted to admins."""
    if not connection.user.admin:
        raise exceptions.NotAuthorizedException


class UserController(Controller):
    path = "/user"

    @post("/login", exclude_from_auth=True)
    async def login(self, data: User, request: Request) -> User:
        """Authenticate user."""
        if data.name in MOCK_DB:
            request.set_session({"user_name": data.name})
            return MOCK_DB[data.name]
        raise exceptions.NotAuthorizedException

    @get("/{user_name:str}")
    async def get_user(self, user_name: str) -> User:
        """Get user by name."""
        if user_name in MOCK_DB:
            return MOCK_DB[user_name]
        raise exceptions.NotFoundException

    @post("/create", guards=[admin_user_guard])
    async def create_user(self, data: User) -> None:
        """Create new user"""
        if data.name in MOCK_DB:
            raise exceptions.HTTPException(status_code=400, detail="User name already taken.")
        MOCK_DB[data.name] = data


@get("/favicon.ico", exclude_from_auth=True)
async def favicon() -> None:
    """Dummy handler to supress error messages in debug mode"""
    ...


async def retrieve_user_handler(session: dict[str, Any], _: "ASGIConnection") -> User:
    """Fetches the current user state from the database and returns it as User instance."""
    user_name = session.get("user_name")
    if user_name in MOCK_DB:
        return MOCK_DB[user_name]
    raise exceptions.NotFoundException

session_config = ServerSideSessionConfig()
session_auth = SessionAuth[User, ServerSideSessionBackend](
    retrieve_user_handler=retrieve_user_handler,
    session_backend_config=session_config,
    # exclude any URLs that should not have authentication.
    exclude=["/schema",],
)


def create_app():
    return Litestar(
        route_handlers=[
            favicon,
            UserController
        ],
        on_app_init=[
            session_auth.on_app_init
        ],
    )


app = create_app()

Test

import asyncio
import sys
from collections.abc import AsyncIterator

import pytest
from litestar import Litestar
from litestar.testing import AsyncTestClient

if sys.platform == "win32":
    asyncio.set_event_loop_policy(asyncio.WindowsSelectorEventLoopPolicy())

pytestmark = pytest.mark.anyio


@pytest.fixture
def anyio_backend() -> str:
    return "asyncio"


@pytest.fixture(name="app")
def fx_app() -> Litestar:
    """App fixture."""
    from app import create_app
    return create_app()


@pytest.fixture(name="client")
async def fx_client(app: Litestar) -> AsyncIterator[AsyncTestClient]:
    """Async client that calls requests on the app.

    Includes all plugins and middleware.
    """
    async with AsyncTestClient(app=app) as client:
        yield client


@pytest.fixture(name="user_client")
async def fx_mission_client() -> AsyncIterator[AsyncTestClient]:
    """Provides a test client for the User Controller"""
    from app import UserController, session_auth, session_config
    from litestar.testing import create_async_test_client

    async with create_async_test_client(
        UserController,
        session_config=session_config,
        on_app_init=[session_auth.on_app_init],
        debug=True
    ) as client:
        await client.set_session_data({"user_name": "admin"})
        yield client


async def test_get_user_not_authorized(client: AsyncTestClient):
    response = await client.get("/user/dummy")
    assert response.status_code == 401


async def test_get_user(user_client: AsyncTestClient):
    response = await user_client.get("/user/dummy")
    assert response.status_code == 200
    data = response.json()
    assert data["name"] == "dummy"


async def test_create_user(user_client: AsyncTestClient):
    new_user = {
        "name": "foo",
        "pw": "bar",
        "admin": False
    }
    response = await user_client.post("/user/create", json=new_user)
    assert response.status_code == 201

[aranvir]

Another thing I thought about is using pytest-mock to override the admin guard within the user_client fixture...

But that lead to really inconsistent behaviour where the mock was applied either for all tests or none (using client or user_client fixture) depending on how I ran the tests (individually, one file, group of files).

mocker.patch("app.admin_user_guard", return_value=None)