Fastapi: [QUESTION] How can I populate data in a test database using the yield session functionality with pytest?

[avocraft]

Description

How can I populate data in a test database dynamically using the yield session functionality with pytest?
What are best practices in FastAPI testing? Is there a good example I could follow?

Additional context

I am getting this error:
AttributeError: 'Depends' object has no attribute 'add'

I have 2 files as an example

test_apps.py

import pytest
from fastapi import Depends

from app.main import app
from config import Config
from data_init import APP
from models.models import AppModel
from test_base import client, get_test_db
from sqlalchemy.orm import Session

@pytest.fixture(scope="session", autouse=True)
def init(db: Session=Depends(get_test_db)):
    new_app = AppModel(**APP)    
    db.add(new_app)
    db.commit()
    db.refresh(new_app)
    AppModel.filter_or_404(db, id=APP["id"])

def test_get_apps(client, db: Session=Depends(get_test_db)):
    response = client.get("/applications")
    assert response.status_code == 200

test_base.py

from typing import Optional, AsyncIterable

import pytest
from fastapi import Depends
from starlette.testclient import TestClient
from sqlalchemy import create_engine
from sqlalchemy.engine import Engine as Database
from sqlalchemy.orm import Session
from sqlalchemy_utils import database_exists, create_database, drop_database

from app.main import app
from app.dependency import get_db
from models.__base import Base
from config import Config

url = str(Config.SQLALCHEMY_DATABASE_URI + "_test")
_db_conn = create_engine(url)


def get_test_db_conn() -> Database:
    assert _db_conn is not None
    return _db_conn


def get_test_db(db_conn=Depends(get_test_db_conn)) -> AsyncIterable[Session]:
    sess = Session(bind=db_conn)

    try:
        yield sess
    finally:
        sess.close()


@pytest.fixture(scope="session", autouse=True)
def create_test_database():
    """
  Create a clean database on every test case.
  For safety, we should abort if a database already exists.

  We use the `sqlalchemy_utils` package here for a few helpers in consistently
  creating and dropping the database.
  """
    try:
        assert not database_exists(url), "Test database already exists. Aborting tests."
        create_database(url)  # Create the test database.
        Base.metadata.create_all(_db_conn)  # Create the tables.
        app.dependency_overrides[get_db] = get_test_db # Mock the Dependency
        yield  # Run the tests.
    finally:
        drop_database(url)  # Drop the test database.


@pytest.fixture()
def client():
    """
    When using the 'client' fixture in test cases, we'll get full database
    rollbacks between test cases:
    """
    with TestClient(app) as client:
        yield client

Is there a better way or a better practice to handle testing?

[dmontagu]

You're very close, you just can't use Depends(...) outside of an endpoint definition. What's happening is that the value of the argument is just the default value, which is of type Depends. This gets changed in endpoint calls by FastAPI, but Depends doesn't get injected into arbitrary calls.

If you remove db: Session = Depends(get_test_db) from the function signatures and instead just put db = get_test_db() as the first line of the function, it should work.

If you want to use dependency injection, in tests, you should just rely on pytest fixtures for it -- so you can just create a pytest fixture that returns the value you are trying to inject via the fastapi Depends function now.

If there is some reason that you don't think you can easily replace the = Depends(get_test_database) with manual function calls, let us know -- it's probably not too hard to work through any remaining issues.

---

In general, I think the approach you are taking looks good for testing, and is very similar to the approach I use in my own projects.

[avocraft]

Thank you for pointing me to the right direction @dmontagu

I got it working like this:

test_base.py

url = str(Config.SQLALCHEMY_DATABASE_URI + "_test")
_db_conn = create_engine(url)

def get_test_db_conn() -> Database:
    assert _db_conn is not None
    return _db_conn


def get_test_db() -> AsyncIterable[Session]:
    sess = Session(bind=_db_conn)

    try:
        yield sess
    finally:
        sess.close()


@pytest.fixture(scope="session", autouse=True)
def create_test_database():
    """
  Create a clean database on every test case.

  We use the `sqlalchemy_utils` package here for a few helpers in consistently
  creating and dropping the database.
  """
    if database_exists(url):
        drop_database(url)
    create_database(url)  # Create the test database.
    Base.metadata.create_all(_db_conn)  # Create the tables.
    app.dependency_overrides[get_db] = get_test_db  # Mock the Database Dependency
    yield  # Run the tests.
    drop_database(url)  # Drop the test database.


@pytest.yield_fixture
def test_db_session():
    """Returns an sqlalchemy session, and after the test tears down everything properly."""

    session = Session(bind=_db_conn)

    yield session
    # Drop all data after each test
    for tbl in reversed(Base.metadata.sorted_tables):
        _db_conn.execute(tbl.delete())
    # put back the connection to the connection pool
    session.close()


@pytest.fixture()
def client():
    """
    When using the 'client' fixture in test cases, we'll get full database
    rollbacks between test cases:
    """
    with TestClient(app) as client:
        yield client

test_app.py

# Required imports 'create_test_database'
from test_base import (
    client,
    get_test_db,
    create_test_database,
    url,
    test_db_session as db,
)


class TestApps:
    def setup(self):
        self.application_url = "/applications"

    @pytest.fixture(autouse=True)
    def setup_db_data(self, db):
        """Set up all the data before each test"""
        new_app = AppModel(**APP)
        db.add(new_app)
        db.commit()
        db.refresh(new_app)

    def test_404(self, client, db):
        response = client.get("/applications")
        assert response.status_code == 200

I am a little confused why create_test_database needs to be imported but never used anywhere. But that will do for now. Thanks !

[dmontagu]

You need to import it otherwise pytest won't know it exists, and won't be able to find it while looking for autouse fixtures.

If you place the autouse fixture in a file called conftest.py that is in the same folder or a parent folder, it will be automatically detected and imported by pytest (conftest.py is like pytest's version of __init__.py). Typically it makes sense to put any autouse=True fixtures in a conftest.py to make sure they get used without needing to import anyting.

[tiangolo]

Thanks @dmontagu ! :bowing_man: :tada:

I guess that solves your use case, right @avocraft ? If so, you can close the issue.

[github-actions[bot]]

Assuming the original issue was solved, it will be automatically closed now. But feel free to add more comments or create new issues.

[duganchen]

I took forever to figure out how to set up a test database (which is mentioned in a "this is left as an exercise to the reader; we have given you all the tools" way in the documentation), and I ended up with a different approach.

The documentation suggests setting up your app like this:

from sqlalchemy import create_engine
from sqlalchemy.orm import sessionmaker

engine = create_engine(
    SQLALCHEMY_DATABASE_URL, connect_args={"check_same_thread": False}
)
SessionLocal = sessionmaker(bind=engine)

def get_db():
    try:
        db = SessionLocal()
        yield db
    finally:
        db.close()

Well, what part of that do you actually want to override for unit tests? If you're setting up a test database, it's SQLALCHEMY_DATABASE_URL. Both Starlet and Pydantic let you set that up to be overridden with environment variables, but that's really designed for deployment and not for unit testing (the envvars are read and applied on import).

My solution? I moved the database URL into the dependency:

def get_db():
    engine = create_engine(
        SQLALCHEMY_DATABASE_URL, connect_args={"check_same_thread": False}
    )
    SessionLocal = sessionmaker(bind=engine)

    try:
        db = SessionLocal()
        yield db
    finally:
        db.close()

And that allows me to override the dependency, which now contains the database connection string, for unit testing:

from pathlib import Path
from tempfile import TemporaryDirectory

def mock_get_db():
    tempdir = TemporaryDirectory()
    path = Path(tempdir.name) / "rdp.sqlite3"
    uri = f"sqlite:///{path}"
    engine = create_engine(uri, connect_args={"check_same_thread": False})
    SessionLocal = sessionmaker(bind=engine)
    models.Base.metadata.create_all(bind=engine)
    try:
        db = SessionLocal()
        yield db
    finally:
        db.close()


app.dependency_overrides[get_db] = mock_get_db

And if I want the state of the test database to persist across multiple client requests:

tempdir = TemporaryDirectory()
path = Path(tempdir.name) / "rdp.sqlite3"
uri = f"sqlite:///{path}"
engine = create_engine(uri, connect_args={"check_same_thread": False})
SessionLocal = sessionmaker(bind=engine)
models.Base.metadata.create_all(bind=engine)

def mock_get_db():
    try:
        db = SessionLocal()
        yield db
    finally:
        db.close()

[timhughes]

@avocraft I ran into problems with this too. The examples given in the docs don't appear (to me) to fully isolate tests. I am sure my way isn't perfect but it is working for me so far and is pretty generic. I built a pretty cutdown app to test it out which may be useful. https://github.com/timhughes/example-fastapi-sqlachemy-pytest/blob/master/tests

[josenava]

Thank you for pointing me to the right direction @dmontagu

I got it working like this:

test_base.py

url = str(Config.SQLALCHEMY_DATABASE_URI + "_test")
_db_conn = create_engine(url)

def get_test_db_conn() -> Database:
    assert _db_conn is not None
    return _db_conn


def get_test_db() -> AsyncIterable[Session]:
    sess = Session(bind=_db_conn)

    try:
        yield sess
    finally:
        sess.close()


@pytest.fixture(scope="session", autouse=True)
def create_test_database():
    """
  Create a clean database on every test case.

  We use the `sqlalchemy_utils` package here for a few helpers in consistently
  creating and dropping the database.
  """
    if database_exists(url):
        drop_database(url)
    create_database(url)  # Create the test database.
    Base.metadata.create_all(_db_conn)  # Create the tables.
    app.dependency_overrides[get_db] = get_test_db  # Mock the Database Dependency
    yield  # Run the tests.
    drop_database(url)  # Drop the test database.


@pytest.yield_fixture
def test_db_session():
    """Returns an sqlalchemy session, and after the test tears down everything properly."""

    session = Session(bind=_db_conn)

    yield session
    # Drop all data after each test
    for tbl in reversed(Base.metadata.sorted_tables):
        _db_conn.execute(tbl.delete())
    # put back the connection to the connection pool
    session.close()


@pytest.fixture()
def client():
    """
    When using the 'client' fixture in test cases, we'll get full database
    rollbacks between test cases:
    """
    with TestClient(app) as client:
        yield client

test_app.py

# Required imports 'create_test_database'
from test_base import (
    client,
    get_test_db,
    create_test_database,
    url,
    test_db_session as db,
)


class TestApps:
    def setup(self):
        self.application_url = "/applications"

    @pytest.fixture(autouse=True)
    def setup_db_data(self, db):
        """Set up all the data before each test"""
        new_app = AppModel(**APP)
        db.add(new_app)
        db.commit()
        db.refresh(new_app)

    def test_404(self, client, db):
        response = client.get("/applications")
        assert response.status_code == 200

I am a little confused why create_test_database needs to be imported but never used anywhere. But that will do for now. Thanks !

Thanks @avocraft and @dmontagu, writing a similar code in the conftest.py worked perfectly. I just changed @pytest.yield_fixture to @pytest.fixture(autouse=True) and a couple of small modifications like using sessionmaker instead of Session.

[lalitvasoya]

@josenava can you give example with a sessionmaker class? I am getting Could not locate a bind configured on mapper mapped class User->user, SQL expression or this Session this error while using the above example in testcases.

[josenava]

Hey @lalitvasoya not sure if this will help but here is how my conftest.py looks like:

SQLALCHEMY_DATABASE_URL = os.environ.get("POSTGRES_URL_TEST", "")

engine = create_engine(SQLALCHEMY_DATABASE_URL)


def get_test_db():
    SessionLocal = sessionmaker(bind=engine)
    test_db = SessionLocal()

    try:
        yield test_db
    finally:
        test_db.close()


@pytest.fixture(scope="session", autouse=True)
def create_test_database():
    """
    Create a clean database on every test case.
    We use the `sqlalchemy_utils` package here for a few helpers in consistently
    creating and dropping the database.
    """
    if database_exists(SQLALCHEMY_DATABASE_URL):
        drop_database(SQLALCHEMY_DATABASE_URL)
    create_database(SQLALCHEMY_DATABASE_URL)  # Create the test database.
    Base.metadata.create_all(engine)  # Create the tables.
    app.dependency_overrides[get_db] = get_test_db  # Mock the Database Dependency
    yield  # Run the tests.
    drop_database(SQLALCHEMY_DATABASE_URL)  # Drop the test database.


@pytest.fixture
def test_db_session():
    """Returns an sqlalchemy session, and after the test tears down everything properly."""

    SessionLocal = sessionmaker(bind=engine)

    session: Session = SessionLocal()
    yield session
    # Drop all data after each test
    for tbl in reversed(Base.metadata.sorted_tables):
        engine.execute(tbl.delete())
    # put back the connection to the connection pool
    session.close()


@pytest.fixture(scope="module")
def client():
    with TestClient(app) as c:
        yield c