fix in is_context_initiated, pop_db_session_from_context, get_db_session_from_context

krylosov-aa · krylosov-aa · commit 22c1e106b887 · 2025-11-30T00:12:35.000+03:00
diff --git a/Makefile b/Makefile
@@ -5,7 +5,7 @@ lint:
 	flake8 .
 
 test:
-	pytest --cov context_async_sqlalchemy examples/fastapi_example/tests examples/starlette_example/tests examples/fastapi_with_pure_asgi_example/tests --cov-report=term-missing
+	pytest --cov context_async_sqlalchemy tests examples/fastapi_example/tests examples/starlette_example/tests examples/fastapi_with_pure_asgi_example/tests --cov-report=term-missing
 
 test_fastapi:
 	pytest examples/fastapi_example/tests
diff --git a/context_async_sqlalchemy/context.py b/context_async_sqlalchemy/context.py
@@ -22,18 +22,18 @@ def is_context_initiated() -> bool:
     """
     Checks whether the context is initiated
     """
-    return bool(_db_session_ctx.get())
+    return _db_session_ctx.get() is not None
 
 
-def pop_db_session_from_context(context_key: str) -> AsyncSession | None:
+def pop_db_session_from_context(connect: DBConnect) -> AsyncSession | None:
     """
     Removes a session from the context
     """
     session_ctx = _db_session_ctx.get()
     if not session_ctx:
         return None
 
-    session: AsyncSession | None = session_ctx.pop(context_key, None)
+    session: AsyncSession | None = session_ctx.pop(connect.context_key, None)
     return session
 
 
@@ -50,12 +50,12 @@ async def reset_db_session_ctx(
     _db_session_ctx.reset(token)
 
 
-def get_db_session_from_context(context_key: str) -> AsyncSession | None:
+def get_db_session_from_context(connect: DBConnect) -> AsyncSession | None:
     """
     Extracts the session from the context
     """
     session_ctx = _get_initiated_context()
-    return session_ctx.get(context_key)
+    return session_ctx.get(connect.context_key)
 
 
 def put_db_session_to_context(
diff --git a/context_async_sqlalchemy/session.py b/context_async_sqlalchemy/session.py
@@ -19,7 +19,7 @@ async def db_session(connect: DBConnect) -> AsyncSession:
         session = await db_session(connect)
         ...
     """
-    session = get_db_session_from_context(connect.context_key)
+    session = get_db_session_from_context(connect)
     if not session:
         session = await connect.create_session()
         put_db_session_to_context(connect, session)
@@ -86,7 +86,7 @@ async def commit_db_session(connect: DBConnect) -> None:
         await your_function_with_db_session()
         await commit_db_session(connect)
     """
-    session = get_db_session_from_context(connect.context_key)
+    session = get_db_session_from_context(connect)
     if session and session.in_transaction():
         await session.commit()
 
@@ -99,7 +99,7 @@ async def rollback_db_session(connect: DBConnect) -> None:
         await your_function_with_db_session()
         await rollback_db_session(connect)
     """
-    session = get_db_session_from_context(connect.context_key)
+    session = get_db_session_from_context(connect)
     if session and session.in_transaction():
         await session.rollback()
 
diff --git a/examples/fastapi_example/routes/atomic_prev_transaction.py b/examples/fastapi_example/routes/atomic_prev_transaction.py
@@ -54,5 +54,5 @@ async def atomic_and_previous_transaction() -> None:
 
 async def _insert_1() -> None:
     session = await db_session(connection)
-    stmt = insert(ExampleTable).values()
+    stmt = insert(ExampleTable)
     await session.execute(stmt)
diff --git a/examples/fastapi_example/routes/auto_rollback_by_status_code.py b/examples/fastapi_example/routes/auto_rollback_by_status_code.py
@@ -12,7 +12,7 @@ async def auto_rollback_by_status_code() -> None:
     let's imagine that an error code was returned.
     """
     session = await db_session(connection)
-    stmt = insert(ExampleTable).values(text="example_with_db_session")
+    stmt = insert(ExampleTable)
     await session.execute(stmt)
 
     raise HTTPException(status_code=500)
diff --git a/examples/fastapi_example/routes/councurrent_queries.py b/examples/fastapi_example/routes/councurrent_queries.py
@@ -31,12 +31,16 @@ async def concurrent_queries() -> None:
         ),
         _insert_non_ctx(),  # new non context session
         _insert_non_ctx_manual(),  # new non context session
+        run_in_new_ctx(  # new context and session with autorollback
+            _insert_with_exception
+        ),
+        return_exceptions=True,
     )
 
 
 async def _insert() -> None:
     session = await db_session(connection)
-    stmt = insert(ExampleTable).values(text="example_multiple_sessions")
+    stmt = insert(ExampleTable)
     await session.execute(stmt)
 
 
@@ -58,7 +62,7 @@ async def _insert_non_ctx() -> None:
     You don't have to use the context to work with sessions at all
     """
     async with new_non_ctx_atomic_session(connection) as session:
-        stmt = insert(ExampleTable).values(text="example_multiple_sessions")
+        stmt = insert(ExampleTable)
         await session.execute(stmt)
 
 
@@ -67,6 +71,13 @@ async def _insert_non_ctx_manual() -> None:
     You don't have to use the context to work with sessions at all
     """
     async with new_non_ctx_session(connection) as session:
-        stmt = insert(ExampleTable).values(text="example_multiple_sessions")
+        stmt = insert(ExampleTable)
         await session.execute(stmt)
         await session.commit()
+
+
+async def _insert_with_exception() -> None:
+    session = await db_session(connection)
+    stmt = insert(ExampleTable)
+    await session.execute(stmt)
+    raise Exception()
diff --git a/examples/fastapi_example/routes/early_connection_close.py b/examples/fastapi_example/routes/early_connection_close.py
@@ -34,7 +34,5 @@ async def _insert() -> None:
         request.
     """
     session = await db_session(connection)
-    stmt = insert(ExampleTable).values(
-        text="example_with_early_connection_close"
-    )
+    stmt = insert(ExampleTable)
     await session.execute(stmt)
diff --git a/examples/fastapi_with_pure_asgi_example/routes/councurrent_queries.py b/examples/fastapi_with_pure_asgi_example/routes/councurrent_queries.py
@@ -31,12 +31,16 @@ async def concurrent_queries() -> None:
         ),
         _insert_non_ctx(),  # new non context session
         _insert_non_ctx_manual(),  # new non context session
+        run_in_new_ctx(  # new context and session with autorollback
+            _insert_with_exception
+        ),
+        return_exceptions=True,
     )
 
 
 async def _insert() -> None:
     session = await db_session(connection)
-    stmt = insert(ExampleTable).values(text="example_multiple_sessions")
+    stmt = insert(ExampleTable)
     await session.execute(stmt)
 
 
@@ -58,7 +62,7 @@ async def _insert_non_ctx() -> None:
     You don't have to use the context to work with sessions at all
     """
     async with new_non_ctx_atomic_session(connection) as session:
-        stmt = insert(ExampleTable).values(text="example_multiple_sessions")
+        stmt = insert(ExampleTable)
         await session.execute(stmt)
 
 
@@ -67,6 +71,13 @@ async def _insert_non_ctx_manual() -> None:
     You don't have to use the context to work with sessions at all
     """
     async with new_non_ctx_session(connection) as session:
-        stmt = insert(ExampleTable).values(text="example_multiple_sessions")
+        stmt = insert(ExampleTable)
         await session.execute(stmt)
         await session.commit()
+
+
+async def _insert_with_exception() -> None:
+    session = await db_session(connection)
+    stmt = insert(ExampleTable)
+    await session.execute(stmt)
+    raise Exception()
diff --git a/examples/starlette_example/routes/councurrent_queries.py b/examples/starlette_example/routes/councurrent_queries.py
@@ -32,13 +32,17 @@ async def concurrent_queries(_: Request) -> JSONResponse:
         ),
         _insert_non_ctx(),  # new non context session
         _insert_non_ctx_manual(),  # new non context session
+        run_in_new_ctx(  # new context and session with autorollback
+            _insert_with_exception
+        ),
+        return_exceptions=True,
     )
     return JSONResponse({})
 
 
 async def _insert() -> None:
     session = await db_session(connection)
-    stmt = insert(ExampleTable).values(text="example_multiple_sessions")
+    stmt = insert(ExampleTable)
     await session.execute(stmt)
 
 
@@ -60,7 +64,7 @@ async def _insert_non_ctx() -> None:
     You don't have to use the context to work with sessions at all
     """
     async with new_non_ctx_atomic_session(connection) as session:
-        stmt = insert(ExampleTable).values(text="example_multiple_sessions")
+        stmt = insert(ExampleTable)
         await session.execute(stmt)
 
 
@@ -69,6 +73,13 @@ async def _insert_non_ctx_manual() -> None:
     You don't have to use the context to work with sessions at all
     """
     async with new_non_ctx_session(connection) as session:
-        stmt = insert(ExampleTable).values(text="example_multiple_sessions")
+        stmt = insert(ExampleTable)
         await session.execute(stmt)
         await session.commit()
+
+
+async def _insert_with_exception() -> None:
+    session = await db_session(connection)
+    stmt = insert(ExampleTable)
+    await session.execute(stmt)
+    raise Exception()
diff --git a/tests/README.md b/tests/README.md
@@ -1,10 +1,9 @@
 # Tests
 
 
-The library is tested in the context of real asynchronous applications.
-For this reason, the library’s tests are essentially the same as the
-example tests -  each example includes a complete test suite that
-verifies all of the library’s capabilities.
+The library is also tested in the context of real asynchronous applications.
+Each example includes a complete test suite that verifies all of the
+library’s capabilities.
 
 - [Tests with FastAPI](../examples/fastapi_example/tests)
 - [Tests with Starlette](../examples/starlette_example/tests)
diff --git a/tests/__init__.py b/tests/__init__.py
diff --git a/tests/conftest.py b/tests/conftest.py
@@ -0,0 +1,18 @@
+"""
+Basic settings and fixtures for testing
+"""
+
+from typing import AsyncGenerator
+
+import pytest_asyncio
+from sqlalchemy.ext.asyncio import AsyncSession
+
+from context_async_sqlalchemy.test_utils import rollback_session
+from examples.database import connection
+
+
+@pytest_asyncio.fixture
+async def db_session_test() -> AsyncGenerator[AsyncSession]:
+    """The session that is used inside the test"""
+    async with rollback_session(connection) as session:
+        yield session
diff --git a/tests/test_context.py b/tests/test_context.py
@@ -0,0 +1,77 @@
+import pytest
+from sqlalchemy.ext.asyncio import AsyncSession
+
+from context_async_sqlalchemy import (
+    get_db_session_from_context,
+    init_db_session_ctx,
+    is_context_initiated,
+    pop_db_session_from_context,
+    put_db_session_to_context,
+    reset_db_session_ctx,
+)
+from examples.database import connection
+
+
+@pytest.mark.asyncio
+async def test_init_db_session_ctx() -> None:
+    assert is_context_initiated() is False
+
+    token = init_db_session_ctx()
+
+    assert is_context_initiated() is True
+
+    try:
+        init_db_session_ctx()
+    except Exception:
+        ...
+    else:
+        raise Exception()
+
+    assert is_context_initiated() is True
+
+    await reset_db_session_ctx(token)
+
+    assert is_context_initiated() is False
+
+
+@pytest.mark.asyncio
+async def test_init_db_session_ctx_force(
+    db_session_test: AsyncSession,
+) -> None:
+    init_db_session_ctx()
+    put_db_session_to_context(connection, db_session_test)
+    assert get_db_session_from_context(connection) is db_session_test
+
+    init_db_session_ctx(force=True)
+    assert get_db_session_from_context(connection) is None
+
+
+@pytest.mark.asyncio
+async def test_pop_db_session_from_context(
+    db_session_test: AsyncSession,
+) -> None:
+    init_db_session_ctx()
+    assert pop_db_session_from_context(connection) is None
+
+    put_db_session_to_context(connection, db_session_test)
+    assert pop_db_session_from_context(connection) is db_session_test
+    assert pop_db_session_from_context(connection) is None
+
+
+@pytest.mark.asyncio
+async def test_get_db_session_from_context(
+    db_session_test: AsyncSession,
+) -> None:
+    try:
+        get_db_session_from_context(connection)
+    except Exception:
+        ...
+    else:
+        raise Exception()
+
+    init_db_session_ctx()
+    assert get_db_session_from_context(connection) is None
+
+    put_db_session_to_context(connection, db_session_test)
+    assert get_db_session_from_context(connection) is db_session_test
+    assert get_db_session_from_context(connection) is db_session_test
