diff --git a/CHANGES.md b/CHANGES.md index 25f8542..a98fbf4 100644 --- a/CHANGES.md +++ b/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 diff --git a/docs/testing.md b/docs/testing.md index 62c391e..e89305a 100644 --- a/docs/testing.md +++ b/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 diff --git a/src/flask_sqlalchemy_lite/_extension.py b/src/flask_sqlalchemy_lite/_extension.py index bcfcc1d..6a6ec85 100644 --- a/src/flask_sqlalchemy_lite/_extension.py +++ b/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 diff --git a/tests/conftest.py b/tests/conftest.py index 49b1646..61ad2a9 100644 --- a/tests/conftest.py +++ b/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) diff --git a/tests/test_isolation.py b/tests/test_isolation.py new file mode 100644 index 0000000..51f2d57 --- /dev/null +++ b/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