Merge pull request #125 from slaarti/test_session_typing

Fix return type for async session fixture in tests.
igorbenav · Jul 9, 2024 · ef29727 · ef29727
2 parents e70ffe5 + ef1a945
commit ef29727
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 2 deletions.
diff --git a/tests/sqlalchemy/conftest.py b/tests/sqlalchemy/conftest.py
@@ -1,3 +1,4 @@
+from collections.abc import AsyncGenerator
 from typing import Optional
 from datetime import datetime
 
@@ -278,7 +279,7 @@ class TaskRead(TaskReadSub):
 
 
 @pytest_asyncio.fixture(scope="function")
-async def async_session() -> AsyncSession:
+async def async_session() -> AsyncGenerator[AsyncSession]:
     session = sessionmaker(async_engine, class_=AsyncSession, expire_on_commit=False)
 
     async with session() as s:

diff --git a/tests/sqlmodel/conftest.py b/tests/sqlmodel/conftest.py
@@ -1,3 +1,4 @@
+from collections.abc import AsyncGenerator
 from typing import Optional
 
 import pytest
@@ -272,7 +273,7 @@ class TaskRead(TaskReadSub):
 
 
 @pytest_asyncio.fixture(scope="function")
-async def async_session() -> AsyncSession:
+async def async_session() -> AsyncGenerator[AsyncSession]:
     session = sessionmaker(async_engine, class_=AsyncSession, expire_on_commit=False)
 
     async with session() as s: