add test isolation manager

CHANGES.md

@@ -5,6 +5,8 @@ Unreleased
 -   Drop support for Python 3.9.
 -   Add `get_or_abort` and `one_or_abort` methods, which get a single row or
     otherwise tell Flask to abort with a 404 error.
+-   Add `test_isolation` context manager, which isolates changes to the database
+    so that tests don't affect each other.
 
 ## Version 0.1.0
 

docs/testing.md

@@ -149,7 +149,8 @@ Each engine in `db.engines` can be patched to represent a connection with a
 transaction instead of a pool. Then all operations will occur inside the
 transaction and be discarded at the end, without writing anything permanently.
 
-Modify the `app` fixture to do this patching.
+Modify the `app` fixture to do this patching with the
+{meth}`.SQLAlchemy.test_isolation` context manager.
 
 ```python
 import pytest
@@ -160,32 +161,17 @@ def app(monkeypatch):
     app = create_app({
         "SQLALCHEMY_ENGINES": {"default": "postgresql:///project-test"}
     })
-    cleanup = []
 
-    with app.app_context():
-        monkeypatch.setitem(
-            db.sessionmaker.kw, "join_transaction_mode", "create_savepoint"
-        )
-
-        for engine in db.engines.values():
-            connection = engine.connect()
-            transaction = connection.begin()
-            cleanup.append(transaction.rollback)
-            cleanup.append(connection.close)
-            connection.close = lambda: None
-            connection.begin = connection.begin_nested
-            monkeypatch.setattr(engine, "connect", lambda _c=connection: _c)
-
-    yield app
-
-    for f in cleanup:
-        f()
+    with db.test_isolation():
+        yield app
 ```
 
 This is not needed when using a SQLite in memory database as discussed above, as
 each test will already be using a separate app with a separate in memory
-database.
+database. If you do use it with SQLite, you'll need to [fix the SQLite driver's
+transaction behavior, as described in SQLAlchemy's docs][transaction].
 
+[transaction]: https://docs.sqlalchemy.org/dialects/sqlite.html#sqlite-transactions
 
 ## Async
 

src/flask_sqlalchemy_lite/_extension.py

@@ -1,7 +1,13 @@
 from __future__ import annotations
 
+import collections.abc as cabc
 import typing as t
+from contextlib import asynccontextmanager
+from contextlib import AsyncExitStack
+from contextlib import contextmanager
+from contextlib import ExitStack
 from dataclasses import dataclass
+from unittest import mock
 from weakref import WeakKeyDictionary
 
 import sqlalchemy as sa
@@ -406,6 +412,89 @@ async def async_one_or_abort(
         except (sa_exc.NoResultFound, sa_exc.MultipleResultsFound):
             abort(code, **(abort_kwargs or {}))
 
+    @contextmanager
+    def test_isolation(self) -> cabc.Iterator[None]:
+        """Context manager to isolate the database during a test. Commits are
+        rolled back when the ``with`` block exits, and will not be seen by other
+        tests.
+
+        This patches the SQLAlchemy engine and session to use a single
+        connection, in a transaction that is rolled back when the context exits.
+
+        If the code being tested uses async features, use
+        :meth:`async_test_isolation` instead. It will isolate both the sync and
+        async operations.
+
+        When using SQLite, follow the `SQLAlchemy docs`__ to fix the driver's
+        transaction handling.
+
+        __ https://docs.sqlalchemy.org/en/20/dialects/sqlite.html#sqlite-transactions
+
+        .. versionadded:: 0.2
+        """
+        with ExitStack() as exit_stack:
+            # Instruct the session to use nested transactions when it sees that
+            # its connection is already in a transaction.
+            exit_stack.enter_context(
+                mock.patch.dict(
+                    self.sessionmaker.kw, {"join_transaction_mode": "create_savepoint"}
+                )
+            )
+
+            for engine in self.engines.values():
+                # Create the connection, to be closed when the context exits.
+                connection: sa.Connection = exit_stack.enter_context(engine.connect())
+                # The connection cannot be closed by code being tested. This
+                # ensures the transaction remains active.
+                connection.close = _nop  # type: ignore[method-assign]
+                # The engine will always return the same connection, with the
+                # active transaction.
+                exit_stack.enter_context(
+                    mock.patch.object(engine, "connect", lambda _c=connection: _c)
+                )
+                # Start the transaction, to be rolled back when the context exits.
+                transaction = connection.begin()
+                exit_stack.callback(transaction.rollback)
+                # If code being tested tries to start the transaction, start a
+                # nested transaction instead.
+                connection.begin = connection.begin_nested  # type: ignore[assignment]
+
+            yield None
+
+    @asynccontextmanager
+    async def async_test_isolation(self) -> cabc.AsyncIterator[None]:
+        """Async version of :meth:`test_isolation` to be used as an
+        ``async with`` block. It will isolate the sync code as well, so you do
+        not need to use ``test_isolation`` as well.
+
+        .. versionadded:: 0.2
+        """
+        async with AsyncExitStack() as exit_stack:
+            # Also isolate the sync operations.
+            exit_stack.enter_context(self.test_isolation())
+
+            await exit_stack.enter_context(
+                mock.patch.dict(
+                    self.async_sessionmaker.kw,
+                    {"join_transaction_mode": "create_savepoint"},
+                )
+            )
+
+            for engine in self.async_engines.values():
+                connection: sa_async.AsyncConnection = (
+                    await exit_stack.enter_async_context(engine.connect())
+                )
+                connection.close = _async_nop  # type: ignore[method-assign]
+                connection.aclose = _async_nop  # type: ignore[method-assign]
+                exit_stack.enter_context(
+                    mock.patch.object(engine, "connect", lambda _c=connection: _c)
+                )
+                transaction = connection.begin()
+                exit_stack.push_async_callback(transaction.rollback)
+                connection.begin = connection.begin_nested  # type: ignore[method-assign]
+
+            yield None
+
 
 @dataclass
 class _State:
@@ -439,3 +528,11 @@ async def _close_async_sessions(e: BaseException | None) -> None:
 
     for session in sessions.values():
         await session.close()
+
+
+def _nop() -> None:
+    pass
+
+
+async def _async_nop() -> None:
+    pass

tests/conftest.py

@@ -2,12 +2,17 @@
 
 import collections.abc as cabc
 import os
+import sys
 import typing as t
 from pathlib import Path
 
 import pytest
+import sqlalchemy as sa
 from flask import Flask
 from flask.ctx import AppContext
+from sqlalchemy import event
+from sqlalchemy.engine.interfaces import DBAPIConnection
+from sqlalchemy.pool import ConnectionPoolEntry
 
 from flask_sqlalchemy_lite import SQLAlchemy
 
@@ -44,4 +49,32 @@ def app_ctx(app: Flask) -> t.Iterator[AppContext]:
 
 @pytest.fixture
 def db(app: Flask) -> SQLAlchemy:
-    return SQLAlchemy(app)
+    engine_options: dict[str, t.Any] = {}
+
+    if sys.version_info >= (3, 12):
+        # Fix sqlite driver's handling of transactions.
+        # https://docs.sqlalchemy.org/en/20/dialects/sqlite.html#sqlite-transactions
+        engine_options["connect_args"] = {"autocommit": False}
+
+    return SQLAlchemy(app, engine_options=engine_options)
+
+
+if sys.version_info < (3, 12):
+    # Fix sqlite3 driver's handling of transactions.
+    # https://docs.sqlalchemy.org/en/20/dialects/sqlite.html#sqlite-transactions
+
+    def _sqlite_connect(
+        dbapi_connection: DBAPIConnection, connection_record: ConnectionPoolEntry
+    ) -> None:
+        dbapi_connection.isolation_level = None
+
+    def _sqlite_begin(conn: sa.Connection) -> None:
+        conn.exec_driver_sql("BEGIN")
+
+    @pytest.fixture(scope="session", autouse=True)
+    def _sqlite_isolation() -> cabc.Iterator[None]:
+        event.listen(sa.Engine, "connect", _sqlite_connect)
+        event.listen(sa.Engine, "begin", _sqlite_begin)
+        yield
+        event.remove(sa.Engine, "begin", _sqlite_begin)
+        event.remove(sa.Engine, "connect", _sqlite_connect)

tests/test_isolation.py

@@ -0,0 +1,41 @@
+from __future__ import annotations
+
+import sqlalchemy as sa
+from flask import Flask
+from sqlalchemy import orm
+
+from flask_sqlalchemy_lite import SQLAlchemy
+
+
+class Base(orm.DeclarativeBase):
+    pass
+
+
+class Todo(Base):
+    __tablename__ = "todo"
+    id: orm.Mapped[int] = orm.mapped_column(primary_key=True)
+
+
+def test_isolation(app: Flask, db: SQLAlchemy) -> None:
+    # Database setup, this item will be present at the start of each isolated block.
+    with app.app_context():
+        Base.metadata.create_all(db.engine)
+        db.session.add(Todo())
+        db.session.commit()
+
+    # Sees the setup item, adds a new item.
+    with app.app_context(), db.test_isolation():
+        db.session.add(Todo())
+        db.session.commit()
+        assert db.session.scalar(sa.select(sa.func.count(Todo.id))) == 2
+
+    # Does not see the previous added item, deletes the setup item.
+    with app.app_context(), db.test_isolation():
+        assert db.session.scalar(sa.select(sa.func.count(Todo.id))) == 1
+        db.session.delete(db.session.get_one(Todo, 1))
+        db.session.commit()
+        assert db.session.scalar(sa.select(sa.func.count(Todo.id))) == 0
+
+    # Deleted setup item has returned.
+    with app.app_context():
+        assert db.session.scalar(sa.select(sa.func.count(Todo.id))) == 1