fix(sqla_factory): added an async context manager in SQLAASyncPersistence (#630)

Author: nisemenov

polyfactory/factories/sqlalchemy_factory.py

@@ -52,16 +52,18 @@ def __init__(self, session: AsyncSession) -> None:
         self.session = session
 
     async def save(self, data: T) -> T:
-        self.session.add(data)
-        await self.session.commit()
-        await self.session.refresh(data)
+        async with self.session as session:
+            session.add(data)
+            await session.commit()
+            await session.refresh(data)
         return data
 
     async def save_many(self, data: list[T]) -> list[T]:
-        self.session.add_all(data)
-        await self.session.commit()
-        for batch_item in data:
-            await self.session.refresh(batch_item)
+        async with self.session as session:
+            session.add_all(data)
+            await session.commit()
+            for batch_item in data:
+                await session.refresh(batch_item)
         return data
 
 

tests/sqlalchemy_factory/test_sqlalchemy_factory_common.py

@@ -18,6 +18,7 @@
     func,
     inspect,
     orm,
+    select,
     text,
     types,
 )
@@ -343,13 +344,15 @@ class Factory(SQLAlchemyFactory[AsyncModel]):
             __async_session__ = session_config(session)
             __model__ = AsyncModel
 
-        result = await Factory.create_async()
-        assert inspect(result).persistent  # type: ignore[union-attr]
+        instance = await Factory.create_async()
+        result = await session.scalar(select(AsyncModel).where(AsyncModel.id == instance.id))
+        assert result
 
         batch_result = await Factory.create_batch_async(size=2)
         assert len(batch_result) == 2
         for batch_item in batch_result:
-            assert inspect(batch_item).persistent  # type: ignore[union-attr]
+            result = await session.scalar(select(AsyncModel).where(AsyncModel.id == batch_item.id))
+            assert result
 
 
 @pytest.mark.parametrize(
@@ -392,8 +395,9 @@ class Factory(SQLAlchemyFactory[AsyncRefreshModel]):
             test_int = Ignore()
             test_bool = Ignore()
 
-        result = await Factory.create_async()
-        assert inspect(result).persistent  # type: ignore[union-attr]
+        instance = await Factory.create_async()
+        result = await session.scalar(select(AsyncRefreshModel).where(AsyncRefreshModel.id == instance.id))
+        assert result
         assert result.test_datetime is not None
         assert isinstance(result.test_datetime, datetime)
         assert result.test_str == "test_str"