How to use PostgreSQL test database in async FastAPI tests?

Question

I'm working on an async FastAPI project and I want to connect to the database during tests. Coming from Django, my instinct was to create pytest fixtures that take care of creating/dropping the test database. However, I couldn't find much documentation on how to do this. The most complete instructions I could find were in this tutorial, but they don't work for me because they are all synchronous. I'm somewhat new to async development so I'm having trouble adapting the code to work async. This is what I have so far:

import pytest
from sqlalchemy.ext.asyncio import create_async_engine, session
from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy.orm import sessionmaker
from sqlalchemy_utils import database_exists, create_database

from fastapi.testclient import TestClient

from app.core.db import get_session
from app.main import app


Base = declarative_base()


@pytest.fixture(scope="session")
def db_engine():
    default_db = (
        "postgresql+asyncpg://postgres:postgres@postgres:5432/postgres"
    )
    test_db = "postgresql+asyncpg://postgres:postgres@postgres:5432/test"
    engine = create_async_engine(default_db)
    if not database_exists(test_db):  # <- Getting error on this line
        create_database(test_db)

    Base.metadata.create_all(bind=engine)
    yield engine


@pytest.fixture(scope="function")
def db(db_engine):
    connection = db_engine.connect()

    # begin a non-ORM transaction
    connection.begin()

    # bind an individual Session to the connection
    Session = sessionmaker(bind=connection)
    db = Session()
    # db = Session(db_engine)

    yield db

    db.rollback()
    connection.close()


@pytest.fixture(scope="function")
def client(db):
    app.dependency_overrides[get_session] = lambda: db

    PREFIX = "/api/v1/my-endpoint"
    with TestClient(PREFIX, app) as c:
        yield c

And this is the error I'm getting:

E           sqlalchemy.exc.MissingGreenlet: greenlet_spawn has not been called; can't call await_() here. Was IO attempted in an unexpected place? (Background on this error at: https://sqlalche.me/e/14/xd2s)

/usr/local/lib/python3.9/site-packages/sqlalchemy/util/_concurrency_py3k.py:67: MissingGreenlet

Any idea what I have to do to fix it?

Answer 1

To anyone looking, here's the solution. The first problem with OP's code is that it's not async, so it should have been:

async def db_engine():

The other - and more significant - problem arises with calling sync code (database_exists(test_db), create_database(test_db), and eventually Base.metadata.create_all(bind=engine)) while using SQLAlchemy's async sessions (which, most frustratingly, uses greenlet and a sync session under the hood). Of course, OP only really reached the first error.

Ideally, you should call the sync code in a threadpool and await that future/coro with something like fastapi.concurrency.run_in_threadpool or asyncio.to_thread. These ideally should work.

But I can tell you firsthand that they don't (and you could try it, maybe you'll get lucky, haha) because greenlet monkey-patches everything - which, have I mentioned, is absolutely infuriating - so you get the same error as before.
I'm not sure if this is because the call is never sent to another thread or because the low-level IO functions will also be monkey-patched the same on another thread.

Instead, you have to wrap your sync code with greenlet directly, so you might have:

from sqlalchemy.util import greenlet_spawn

@pytest.fixture(scope="session")
async def db_engine():
    default_db = (
         "postgresql+asyncpg://postgres:postgres@postgres:5432/postgres"
    )
    test_db = "postgresql+asyncpg://postgres:postgres@postgres:5432/test"
    engine = create_async_engine(default_db)

    def _make_test_db():
        if not database_exists(test_db):
            create_database(test_db)

    def _make_test_tables():
        Base.metadata.create_all(bind=engine)

    await greenlet_spawn(_make_test_db)
    await greenlet_spawn(_make_test_tables)
    yield engine

My use case also had a drop test db like

def _make_test_db():
    if database_exists(TEST_DATABSE_URL):
        drop_database(TEST_DATABSE_URL)
    create_database(TEST_DATABSE_URL)

def _drop_test_db():
    drop_database(TEST_DATABSE_URL)

await greenlet_spawn(_make_test_db)
yield
await greenlet_spawn(_drop_test_db)

I wasted a lot of time on this, but I've decided to abandon async sqlalchemy entirely and simply use sync for db, routes, tests, everything.
I'd suggest you consider the same because if, for example, fastapi.concurrency.run_in_threadpool doesn't work, what happens when you work with an actual sync library and fastapi tries to run it, will you have to replace fastapi.concurrency.run_in_threadpool with sqlalchemy.util.greenlet_spawn? And every other library that must run sync code under the hood?

As a closing note, there are other things wrong with OP's code such as making the other fixtures async, using async_sessionmaker instead of sessionmaker, using an async test client instead of the sync TestClient from fastapi (which is from starlette, which is from httpx).
Other, less relevant, suggestions might include subclassing sqlalchemy.orm.DeclarativeBase (instead of using declarative_base) to declare Base, and moving the dependency overrides to another, session-scoped, fixture.