How to rollback database changes in tests (SQLAlchemy)?

[orsinium]

I'm using the built-in litestar plugin for SQLAlchemy and I'm trying to figure out how to make it use the same connection as the one I created in tests. The goal is to run a rollback on the connection after each test to avoid side effects. In my current implementation, I get the following error:

RuntimeError: Task <Task pending name='anyio.from_thread.BlockingPortal._call_func' coro=<BlockingPortal._call_func() running at /.../anyio/from_thread.py:219> cb=[TaskGroup._spawn.<locals>.task_done() at /.../anyio/_backends/_asyncio.py:726]> got Future <Future pending cb=[Protocol._on_waiter_completed()]> attached to a different loop

My guess is that litestar runs in a different event loop from the one that is used by pytest-asyncio but I don't see a way to pass an event loop into the app. Did anyone already solve the problem? Do you have rollbacks in your tests?

This is the full code of the test fixtures:

import pytest
import pytest_asyncio
from litestar import Litestar
from litestar.testing import AsyncTestClient
from sqlalchemy.ext.asyncio import AsyncConnection, AsyncEngine, AsyncSession, create_async_engine

from project._app import create_app
from project._settings import Settings

@pytest_asyncio.fixture
async def db_engine() -> AsyncEngine:
    settings = Settings()
    return create_async_engine(settings.database_url)

@pytest_asyncio.fixture
async def db_conn(db_engine: AsyncEngine):
    async with db_engine.connect() as conn:
        async with conn.begin() as transaction:
            yield conn
            await transaction.rollback()

@pytest.fixture
def app(db_conn: AsyncConnection, db_engine: AsyncEngine) -> Litestar:
    return create_app(
        session_maker=lambda: AsyncSession(bind=db_conn),
        engine_instance=db_engine,
    )

@pytest_asyncio.fixture
async def client(app: Litestar):
    async with AsyncTestClient(app) as client:
        yield client

[JacobCoffee]

@cofin

[orsinium]

I couldn't figure out a solution that didn't fall apart and moved to fastapi. See #1920

[provinzkraut]

If I understand your comment there correctly you didn’t actually solve this by migrating to FastAPI, and you experienced the same problem there, but rather by using httpx.AsyncClient. Would you mush sharing some details about that?

[orsinium]

Both sync and async clients in litestar spawn a new thread that runs the asyncio event loop with the application. That means, you cannot use within application jobs that you started in your tests because there are created in different threads (and so in different even loops). This behavior, supposedly, is copied in litestar from the sync httpx client, and so you have the same problem in fastapi when using sync httpx client. To solve the issue, you either need to use async httpx client or to use sync sqlalchemy session. The latter wasn't an option for me because that means your whole application cannot use async sqlalchemy API anymore.

[furioness]

A Django-like approach: wrap everything into a global transaction per app's lifecycle.

Haven't tested on any complicated code yet, but at least it looks like a good approach.

from contextlib import asynccontextmanager
from typing import AsyncGenerator, cast

from litestar import Litestar
from litestar.datastructures import State
from litestar.plugins.sqlalchemy import SQLAlchemyAsyncConfig
from litestar.types import Scope
from sqlalchemy import event
from sqlalchemy.ext.asyncio import AsyncEngine, AsyncSession, async_sessionmaker


class SQLAlchemyAsyncConfigTransactional(SQLAlchemyAsyncConfig):
    """
    Wraps app's lifecycle into one transaction.
    So multiple commits/rollbacks per single test are possible.

    - One engine connection plus one outer transaction for the app lifespan (per test)
    - Sessionmaker bound to that connection
    - Each provided session runs inside a SAVEPOINT that is recreated after each commit/rollback
    """

    @asynccontextmanager
    async def lifespan(self, app: Litestar) -> AsyncGenerator[None, None]:
        deps = self.create_app_state_items()
        app.state.update(deps)

        engine = cast(AsyncEngine, app.state[self.engine_app_state_key])
        conn = await engine.connect()
        outer_tx = await conn.begin()

        # Replace the sessionmaker with one bound to this single connection
        session_maker = async_sessionmaker(bind=conn,
            expire_on_commit=getattr(self.session_config, "expire_on_commit", False),
            class_=AsyncSession, )
        app.state[self.session_maker_app_state_key] = session_maker
        try:
            if self.create_all:
                await self.create_all_metadata(app)
            yield
        finally:
            # Roll back everything done during the lifespan (i.e., per test)
            if outer_tx.is_active:
                await outer_tx.rollback()

            await conn.close()

            # Dispose engine created by base config
            eng = app.state.get(self.engine_app_state_key)
            if eng is not None and hasattr(eng, "dispose"):
                await cast(AsyncEngine, eng).dispose()

    def provide_session(self, state: State, scope: Scope) -> AsyncSession:
        """
        Create a session instance that would be rolled back on app's shutdown.
        """

        session = super().provide_session(state, scope)

        # Expected only 1 session per lifecycle
        # in case anyone called this method more than once
        if "_savepoint_listener_installed" not in session.info:
            session.sync_session.begin_nested()

            @event.listens_for(session.sync_session, "after_transaction_end")
            def _restart_savepoint(sess, trans):
                if trans.nested and not trans._parent.is_active:
                    sess.begin_nested()

            session.info["_savepoint_listener_installed"] = True

        return session

Use like this

sqlalchemy_config = SQLAlchemyAsyncConfigTransactional(
        connection_string=DB_URI, session_config=session_config, create_all=False)
app = Litestar(plugins=[SQLAlchemyPlugin(config=sqlalchemy_config)])

P.S.
Yes, AI was used to help here.