feat(pytest): support startup in tests (#299)

Author: arththebird

fastapi_sqla/_pytest_plugin.py

@@ -1,16 +1,23 @@
 import os
+from collections.abc import AsyncGenerator, Generator
 from unittest.mock import patch
 from urllib.parse import urlsplit, urlunsplit
 
 from alembic import command
 from alembic.config import Config
-from pytest import fixture
+from pytest import FixtureRequest, fixture
 from sqlalchemy import create_engine, text
-from sqlalchemy.orm.session import sessionmaker
+from sqlalchemy.engine import Connection, Engine
+from sqlalchemy.orm.session import Session, sessionmaker
 
 try:
     import asyncpg  # noqa
-    from sqlalchemy.ext.asyncio import create_async_engine
+    from sqlalchemy.ext.asyncio import (
+        create_async_engine,
+        AsyncEngine,
+        AsyncConnection,
+        AsyncSession,
+    )
 
     asyncio_support = True
 except ImportError:
@@ -22,7 +29,7 @@ def pytest_configure(config):
 
 
 @fixture(scope="session")
-def db_host():
+def db_host() -> str:
     """Default db host used by depending fixtures.
 
     When CI key is set in environment variables, it uses `postgres` as host name else,
@@ -32,7 +39,7 @@ def db_host():
 
 
 @fixture(scope="session")
-def db_user():
+def db_user() -> str:
     """Default db user used by depending fixtures.
 
     postgres
@@ -41,7 +48,7 @@ def db_user():
 
 
 @fixture(scope="session")
-def db_url(db_host, db_user):
+def db_url(db_host: str, db_user: str) -> str:
     """Default db url used by depending fixtures.
 
     db url example postgresql://{db_user}@{db_host}/postgres
@@ -50,24 +57,24 @@ def db_url(db_host, db_user):
 
 
 @fixture(scope="session")
-def engine(db_url):
+def engine(db_url: str) -> Engine:
     return create_engine(db_url)
 
 
 @fixture(scope="session")
-def sqla_connection(engine):
+def sqla_connection(engine: Engine) -> Generator[Connection]:
     with engine.connect() as connection:
         yield connection
 
 
 @fixture(scope="session")
-def alembic_ini_path():  # pragma: no cover
+def alembic_ini_path() -> str:  # pragma: no cover
     """Path for alembic.ini file, defaults to `./alembic.ini`."""
     return "./alembic.ini"
 
 
 @fixture(scope="session")
-def db_migration(db_url, sqla_connection, alembic_ini_path):
+def db_migration(db_url: str, sqla_connection: Connection, alembic_ini_path: str):
     """Run alembic upgrade at test session setup and downgrade at tear down.
 
     Override fixture `alembic_ini_path` to change path of `alembic.ini` file.
@@ -94,54 +101,52 @@ def sqla_modules():
 
 
 @fixture
-def sqla_reflection(sqla_modules, sqla_connection):
+def sqla_reflection(sqla_modules, sqla_connection: Connection):
     import fastapi_sqla
 
-    fastapi_sqla.Base.metadata.bind = sqla_connection
+    fastapi_sqla.Base.metadata.bind = sqla_connection  # type: ignore
     fastapi_sqla.Base.prepare(sqla_connection.engine)
 
 
 @fixture
-def patch_engine_from_config(request, sqla_connection):
+def patch_new_engine(request: FixtureRequest, sqla_connection: Connection):
     """So that all DB operations are never written to db for real."""
     if "dont_patch_engines" in request.keywords:
         yield
     else:
-        transaction = sqla_connection.begin()
-
-        with patch("fastapi_sqla.sqla.engine_from_config") as engine_from_config:
-            engine_from_config.return_value = sqla_connection
-            yield
+        with sqla_connection.begin() as transaction:
+            with patch("fastapi_sqla.sqla.new_engine", return_value=sqla_connection):
+                yield
 
-        transaction.rollback()
+            transaction.rollback()
 
 
 @fixture
-def session_factory():
-    return sessionmaker()
+def session_factory(
+    sqla_connection: Connection, sqla_reflection, patch_new_engine
+) -> sessionmaker:
+    return sessionmaker(bind=sqla_connection)
 
 
 @fixture
-def session(
-    session_factory, sqla_connection, sqla_reflection, patch_engine_from_config
-):
+def session(session_factory: sessionmaker) -> Generator[Session]:
     """Sqla session to use when creating db fixtures.
 
     While it does not write any record in DB, the application will still be able to
     access any record committed with that session.
     """
-    session = session_factory(bind=sqla_connection)
+    session: Session = session_factory()
     yield session
     session.close()
 
 
-def format_async_async_sqlalchemy_url(url):
+def format_async_async_sqlalchemy_url(url: str) -> str:
     scheme, location, path, query, fragment = urlsplit(url)
     return urlunsplit([f"{scheme}+asyncpg", location, path, query, fragment])
 
 
 @fixture(scope="session")
-def async_sqlalchemy_url(db_url):
+def async_sqlalchemy_url(db_url: str) -> str:
     """Default async db url.
 
     It is the same as `db_url` with `postgresql+asyncpg://` as scheme.
@@ -152,46 +157,56 @@ def async_sqlalchemy_url(db_url):
 if asyncio_support:
 
     @fixture
-    def async_engine(async_sqlalchemy_url):
+    def async_engine(async_sqlalchemy_url: str) -> AsyncEngine:
         return create_async_engine(async_sqlalchemy_url)
 
     @fixture
-    async def async_sqla_connection(async_engine):
+    async def async_sqla_connection(
+        async_engine: AsyncEngine,
+    ) -> AsyncGenerator[AsyncConnection]:
         async with async_engine.connect() as connection:
             yield connection
 
     @fixture
-    async def patch_new_engine(request, async_sqla_connection):
+    async def patch_new_async_engine(
+        request: FixtureRequest, async_sqla_connection: AsyncConnection
+    ):
         """So that all async DB operations are never written to db for real."""
         if "dont_patch_engines" in request.keywords:
             yield
         else:
             async with async_sqla_connection.begin() as transaction:
-                with patch("fastapi_sqla.async_sqla.new_engine") as new_engine:
-                    new_engine.return_value = async_sqla_connection
+                with patch(
+                    "fastapi_sqla.async_sqla.new_async_engine",
+                    return_value=async_sqla_connection,
+                ):
                     yield
 
                 await transaction.rollback()
 
     @fixture
-    async def async_sqla_reflection(sqla_modules, async_sqla_connection):
+    async def async_sqla_reflection(
+        sqla_modules, async_sqla_connection: AsyncConnection
+    ):
         from fastapi_sqla import Base
 
         await async_sqla_connection.run_sync(lambda conn: Base.prepare(conn.engine))
 
     @fixture
-    def async_session_factory():
-        from fastapi_sqla.async_sqla import SqlaAsyncSession
-
-        return sessionmaker(class_=SqlaAsyncSession)
+    def async_session_factory(
+        async_sqla_connection: AsyncConnection,
+        async_sqla_reflection,
+        patch_new_async_engine,
+    ) -> sessionmaker:
+        # TODO: Use async_sessionmaker once only supporting 2.x+
+        return sessionmaker(
+            bind=async_sqla_connection, expire_on_commit=False, class_=AsyncSession
+        )  # type: ignore
 
     @fixture
     async def async_session(
-        async_session_factory,
-        async_sqla_connection,
-        async_sqla_reflection,
-        patch_new_engine,
-    ):
-        session = async_session_factory(bind=async_sqla_connection)
+        async_session_factory: sessionmaker,
+    ) -> AsyncGenerator[AsyncSession]:
+        session: AsyncSession = async_session_factory()
         yield session
         await session.close()

fastapi_sqla/async_sqla.py

@@ -1,38 +1,53 @@
+import os
 from collections.abc import AsyncGenerator
 from contextlib import asynccontextmanager
-from typing import Annotated
+from typing import Annotated, Union
 
 import structlog
 from fastapi import Depends, Request, Response
 from fastapi.responses import PlainTextResponse
 from sqlalchemy import text
-from sqlalchemy.ext.asyncio import AsyncEngine
+from sqlalchemy.ext.asyncio import (
+    AsyncConnection,
+    AsyncEngine,
+    async_engine_from_config,
+)
 from sqlalchemy.ext.asyncio import AsyncSession as SqlaAsyncSession
 from sqlalchemy.orm.session import sessionmaker
 from starlette.types import ASGIApp, Message, Receive, Scope, Send
 
 from fastapi_sqla import aws_aurora_support, aws_rds_iam_support
-from fastapi_sqla.sqla import _DEFAULT_SESSION_KEY, Base, new_engine
+from fastapi_sqla.sqla import _DEFAULT_SESSION_KEY, Base, get_envvar_prefix
 
 logger = structlog.get_logger(__name__)
 
 _ASYNC_REQUEST_SESSION_KEY = "fastapi_sqla_async_session"
 _async_session_factories: dict[str, sessionmaker] = {}
 
 
-def new_async_engine(key: str = _DEFAULT_SESSION_KEY):
-    engine = new_engine(key)
-    return AsyncEngine(engine)
+def new_async_engine(
+    key: str = _DEFAULT_SESSION_KEY,
+) -> Union[AsyncEngine, AsyncConnection]:
+    envvar_prefix = get_envvar_prefix(key)
+    lowercase_environ = {k.lower(): v for k, v in os.environ.items()}
+    lowercase_environ.pop(f"{envvar_prefix}warn_20", None)
+    return async_engine_from_config(lowercase_environ, prefix=envvar_prefix)
 
 
 async def startup(key: str = _DEFAULT_SESSION_KEY):
-    engine = new_async_engine(key)
-    aws_rds_iam_support.setup(engine.sync_engine)
-    aws_aurora_support.setup(engine.sync_engine)
+    engine_or_connection = new_async_engine(key)
+    aws_rds_iam_support.setup(engine_or_connection.sync_engine)
+    aws_aurora_support.setup(engine_or_connection.sync_engine)
+
+    async_engine = (
+        engine_or_connection
+        if isinstance(engine_or_connection, AsyncEngine)
+        else engine_or_connection.engine
+    )
 
     # Fail early
     try:
-        async with engine.connect() as connection:
+        async with async_engine.connect() as connection:
             await connection.execute(text("select 'ok'"))
     except Exception:
         logger.critical(
@@ -41,14 +56,15 @@ async def startup(key: str = _DEFAULT_SESSION_KEY):
         )
         raise
 
-    async with engine.connect() as connection:
+    async with async_engine.connect() as connection:
         await connection.run_sync(lambda conn: Base.prepare(conn.engine))
 
+    # TODO: Use async_sessionmaker once only supporting 2.x+
     _async_session_factories[key] = sessionmaker(
-        class_=SqlaAsyncSession, bind=engine, expire_on_commit=False
-    )
+        class_=SqlaAsyncSession, bind=engine_or_connection, expire_on_commit=False
+    )  # type: ignore
 
-    logger.info("engine startup", engine_key=key, async_engine=engine)
+    logger.info("engine startup", engine_key=key, async_engine=engine_or_connection)
 
 
 @asynccontextmanager

fastapi_sqla/base.py

@@ -1,10 +1,11 @@
 import functools
 import os
 import re
+from typing import Union
 
 from deprecated import deprecated
 from fastapi import FastAPI
-from sqlalchemy.engine import Engine
+from sqlalchemy.engine import Connection, Engine
 
 from fastapi_sqla import sqla
 
@@ -72,5 +73,5 @@ def _get_engine_keys() -> set[str]:
     return keys
 
 
-def _is_async_dialect(engine: Engine):
+def _is_async_dialect(engine: Union[Engine, Connection]):
     return engine.dialect.is_async if hasattr(engine.dialect, "is_async") else False

fastapi_sqla/sqla.py

@@ -2,14 +2,14 @@
 import os
 from collections.abc import Generator
 from contextlib import contextmanager
-from typing import Annotated
+from typing import Annotated, Union
 
 import structlog
 from fastapi import Depends, Request, Response
 from fastapi.concurrency import contextmanager_in_threadpool
 from fastapi.responses import PlainTextResponse
 from sqlalchemy import engine_from_config, text
-from sqlalchemy.engine import Engine
+from sqlalchemy.engine import Connection, Engine
 from sqlalchemy.ext.declarative import DeferredReflection
 from sqlalchemy.orm.session import Session as SqlaSession
 from sqlalchemy.orm.session import sessionmaker
@@ -42,24 +42,29 @@ class Base(DeclarativeBase, DeferredReflection):
     __abstract__ = True
 
 
-def new_engine(key: str = _DEFAULT_SESSION_KEY) -> Engine:
+def get_envvar_prefix(key: str) -> str:
     envvar_prefix = "sqlalchemy_"
     if key != _DEFAULT_SESSION_KEY:
         envvar_prefix = f"fastapi_sqla__{key}__{envvar_prefix}"
 
+    return envvar_prefix
+
+
+def new_engine(key: str = _DEFAULT_SESSION_KEY) -> Union[Engine, Connection]:
+    envvar_prefix = get_envvar_prefix(key)
     lowercase_environ = {k.lower(): v for k, v in os.environ.items()}
     lowercase_environ.pop(f"{envvar_prefix}warn_20", None)
     return engine_from_config(lowercase_environ, prefix=envvar_prefix)
 
 
 def startup(key: str = _DEFAULT_SESSION_KEY):
-    engine = new_engine(key)
-    aws_rds_iam_support.setup(engine.engine)
-    aws_aurora_support.setup(engine.engine)
+    engine_or_connection = new_engine(key)
+    aws_rds_iam_support.setup(engine_or_connection.engine)
+    aws_aurora_support.setup(engine_or_connection.engine)
 
     # Fail early
     try:
-        with engine.connect() as connection:
+        with engine_or_connection.engine.connect() as connection:
             connection.execute(text("select 'OK'"))
     except Exception:
         logger.critical(
@@ -68,11 +73,13 @@ def startup(key: str = _DEFAULT_SESSION_KEY):
         )
         raise
 
-    Base.prepare(engine)
+    Base.prepare(engine_or_connection.engine)
 
-    _session_factories[key] = sessionmaker(bind=engine, class_=SqlaSession)
+    _session_factories[key] = sessionmaker(
+        bind=engine_or_connection, class_=SqlaSession
+    )
 
-    logger.info("engine startup", engine_key=key, engine=engine)
+    logger.info("engine startup", engine_key=key, engine=engine_or_connection)
 
 
 @contextmanager