test utils

Author: krylosov-aa

README.md

@@ -265,3 +265,19 @@ replica = DBConnect(
     before_create_session_handler=renew_replica_connect,
 )
 ```
+
+## Testing
+
+The library provides several ready-made utils that can be used in tests,
+for example in fixtures. It helps write tests that share a common transaction
+between the test and the application, so data isolation between tests is
+achieved through fast transaction rollback.
+
+
+You can see the capabilities in the examples:
+
+[Here are tests with a common transaction between the
+application and the tests.](https://github.com/krylosov-aa/context-async-sqlalchemy/blob/main/exmaples/fastapi_example/tests/transactional/__init__.py)
+
+
+[And here's an example with different transactions.](https://github.com/krylosov-aa/context-async-sqlalchemy/blob/main/exmaples/fastapi_example/tests/non_transactional/__init__.py)

context_async_sqlalchemy/context.py

@@ -3,6 +3,8 @@
 
 from sqlalchemy.ext.asyncio import AsyncSession
 
+from .connect import DBConnect
+
 
 def init_db_session_ctx() -> Token[dict[str, AsyncSession] | None]:
     """
@@ -55,14 +57,14 @@ def get_db_session_from_context(context_key: str) -> AsyncSession | None:
 
 
 def put_db_session_to_context(
-    context_key: str,
+    connection: DBConnect,
     session: AsyncSession,
 ) -> None:
     """
     Puts the session into context
     """
     session_ctx = _get_initiated_context()
-    session_ctx[context_key] = session
+    session_ctx[connection.context_key] = session
 
 
 def sessions_stream() -> Generator[AsyncSession, Any, None]:

context_async_sqlalchemy/session.py

@@ -22,7 +22,7 @@ async def db_session(connect: DBConnect) -> AsyncSession:
     session = get_db_session_from_context(connect.context_key)
     if not session:
         session = await connect.create_session()
-        put_db_session_to_context(connect.context_key, session)
+        put_db_session_to_context(connect, session)
     return session
 
 

context_async_sqlalchemy/test_utils.py

@@ -0,0 +1,66 @@
+"""Utilities to use during testing"""
+
+from contextlib import asynccontextmanager
+from typing import AsyncGenerator
+
+from sqlalchemy.ext.asyncio import AsyncSession
+
+from context_async_sqlalchemy import (
+    DBConnect,
+    init_db_session_ctx,
+    put_db_session_to_context,
+    reset_db_session_ctx,
+)
+
+
+@asynccontextmanager
+async def rollback_session(
+    connection: DBConnect,
+) -> AsyncGenerator[AsyncSession]:
+    """A session that always rolls back"""
+    session_maker = await connection.get_session_maker()
+    async with session_maker() as session:
+        try:
+            yield session
+        finally:
+            await session.rollback()
+
+
+@asynccontextmanager
+async def set_test_context() -> AsyncGenerator[None]:
+    """
+    Opens a context similar to middleware, but doesn't commit or
+        rollback automatically. This task falls to the fixture in tests.
+    """
+    token = init_db_session_ctx()
+    try:
+        yield
+    finally:
+        await reset_db_session_ctx(
+            token,
+            # Don't close the session here, as you opened in fixture.
+            with_close=False,
+        )
+
+
+@asynccontextmanager
+async def put_savepoint_session_in_ctx(
+    connection: DBConnect,
+    session: AsyncSession,
+) -> AsyncGenerator[None]:
+    """
+    Sets the context to a session that uses a save point instead of creating
+        a transaction. You need to pass the session you're using inside
+        your tests to attach a new session to the same connection.
+
+    It is also important to use this function inside set_test_context.
+    """
+    session_maker = await connection.get_session_maker()
+    async with session_maker(
+        # Bind to the same connection
+        bind=await session.connection(),
+        # Instead of opening a transaction, it creates a save point
+        join_transaction_mode="create_savepoint",
+    ) as session:
+        put_db_session_to_context(connection, session)
+    yield

exmaples/fastapi_example/tests/conftest.py

@@ -7,8 +7,9 @@
 import pytest_asyncio
 from fastapi import FastAPI
 from httpx import AsyncClient, ASGITransport
-from sqlalchemy.ext.asyncio import async_sessionmaker, AsyncSession
+from sqlalchemy.ext.asyncio import AsyncSession
 
+from context_async_sqlalchemy.test_utils import rollback_session
 from exmaples.fastapi_example.database import master
 from exmaples.fastapi_example.setup_app import lifespan, setup_app
 
@@ -34,20 +35,7 @@ async def client(app: FastAPI) -> AsyncGenerator[AsyncClient]:
 
 
 @pytest_asyncio.fixture
-async def db_session_test(
-    session_maker_test: async_sessionmaker[AsyncSession],
-) -> AsyncGenerator[AsyncSession]:
+async def db_session_test() -> AsyncGenerator[AsyncSession]:
     """The session that is used inside the test"""
-    async with session_maker_test() as session:
-        try:
-            yield session
-        finally:
-            await session.rollback()
-
-
-@pytest_asyncio.fixture
-async def session_maker_test(
-    app: FastAPI,  # To make the connection to the database in lifespan
-) -> AsyncGenerator[async_sessionmaker[AsyncSession]]:
-    session_maker = await master.get_session_maker()
-    yield session_maker
+    async with rollback_session(master) as session:
+        yield session

exmaples/fastapi_example/tests/transactional/conftest.py

@@ -13,15 +13,13 @@
 from typing import AsyncGenerator
 
 import pytest_asyncio
-
 from sqlalchemy.ext.asyncio import AsyncSession
 
-from exmaples.fastapi_example.database import master
-from context_async_sqlalchemy import (
-    init_db_session_ctx,
-    put_db_session_to_context,
-    reset_db_session_ctx,
+from context_async_sqlalchemy.test_utils import (
+    put_savepoint_session_in_ctx,
+    set_test_context,
 )
+from exmaples.fastapi_example.database import master
 
 
 @pytest_asyncio.fixture(autouse=True)
@@ -33,21 +31,6 @@ async def db_session_override(
     The middleware has a special check that won't initialize the context
         if it already exists.
     """
-    token = init_db_session_ctx()
-
-    # Here we create a new session with save point behavior.
-    # This means that committing within the application will save and
-    #   release the save point, rather than committing the entire transaction.
-    conn = await db_session_test.connection()
-    new_session = AsyncSession(
-        bind=conn, join_transaction_mode="create_savepoint"
-    )
-    put_db_session_to_context(master.context_key, new_session)
-    try:
-        yield
-    finally:
-        await reset_db_session_ctx(
-            token,
-            # Don't close the session here, as you opened in fixture.
-            with_close=False,
-        )
+    async with set_test_context():
+        async with put_savepoint_session_in_ctx(master, db_session_test):
+            yield

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "context-async-sqlalchemy"
-version = "1.2.1"
+version = "1.2.2"
 description = "A convenient way to configure and interact with a async sqlalchemy session through context in asynchronous applications"
 readme = "README.md"
 authors = [