Rollback when using pytest

[DanielAndreasen]

First Check

• I added a very descriptive title here.
• I used the GitHub search to find a similar question and didn't find it.
• I searched the SQLModel documentation, with the integrated search.
• I already searched in Google "How to X in SQLModel" and didn't find any information.
• I already read and followed all the tutorial in the docs and didn't find an answer.
• I already checked if it is not related to SQLModel but to Pydantic.
• I already checked if it is not related to SQLModel but to SQLAlchemy.

Commit to Help

• I commit to help with one of those options 👆

Example Code

import os

import pytest
import sqlalchemy as sa
from sqlalchemy.orm import sessionmaker
from sqlmodel import Session

url = f"postgresql+psycopg2://{os.getenv('PGUSER')}:{os.getenv('PGPASSWORD')}@{os.getenv('PGHOST')}/{os.getenv('MYDB')}"
engine = sa.create_engine(url)
TestingSessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)


@pytest.fixture()
def session():
    connection = engine.connect()
    transaction = connection.begin()
    session = TestingSessionLocal(bind=connection)

    nested = connection.begin_nested()

    @sa.event.listens_for(session, "after_transaction_end")
    def end_savepoint(session, transaction):
        nonlocal nested
        if not nested.is_active:
            nested = connection.begin_nested()

    yield session

    # Run this after the test (regardless of status)
    session.close()
    transaction.rollback()
    connection.close()

Description

I want to test the API I'm writing with fastAPI and SQLModel. The API commits to the database, but I want to make a rollback when testing.
The code attached shows how to do this without problem using sqlalchemy.

The problem is that the session from sqlalchemy is slighty different than the one from SQLModel, which I prefer to use, as this is actually the one being used in the API.

If I change the line:
session = TestingSessionLocal(bind=connection)
to:
sesion = Session(engine) # Session from SQLModel
then it does not enter the sa.event.listens_for, and the rollback never happens.

I'm open for another solution as well.

Operating System

Linux

Operating System Details

No response

SQLModel Version

0.0.16

Python Version

Python 3.12.0

Additional Context

Inspiration found here.
https://stackoverflow.com/questions/67255653/how-to-set-up-and-tear-down-a-database-between-
tests-in-fastapi

[DanielAndreasen]

I was able to solve this after many trials and errors (mostly errors).

This is the code:

import pytest
import sqlalchemy as sa
from sqlmodel import Session

from myapp.db.database import engine


@pytest.fixture()
def session():
    connection = engine.connect()
    transaction = connection.begin()
    session = Session(bind=connection)

    nested = connection.begin_nested()

    @sa.event.listens_for(session, "after_transaction_end")
    def end_savepoint(session, transaction):
        nonlocal nested
        if not nested.is_active:
            nested = connection.begin_nested()

    yield session

    session.close()
    transaction.rollback()
    connection.close()

The trick was to use Session(bind=connection) instead of Session(engine)