mirror of
https://github.com/d3vyce/fastapi-toolsets.git
synced 2026-04-16 06:36:26 +02:00
feat: move db related function from pytest to db module (#119)
This commit is contained in:
d3vyce
GitHub
@@ -1,12 +1,12 @@
|
||||
"""Pytest helper utilities for FastAPI testing."""
|
||||
|
||||
import os
|
||||
import warnings
|
||||
from collections.abc import AsyncGenerator, Callable
|
||||
from contextlib import asynccontextmanager
|
||||
from typing import Any
|
||||
|
||||
from httpx import ASGITransport, AsyncClient
|
||||
from sqlalchemy import text
|
||||
from sqlalchemy.engine import make_url
|
||||
from sqlalchemy.ext.asyncio import (
|
||||
AsyncSession,
|
||||
@@ -15,7 +15,134 @@ from sqlalchemy.ext.asyncio import (
|
||||
)
|
||||
from sqlalchemy.orm import DeclarativeBase
|
||||
|
||||
from ..db import create_db_context
|
||||
from sqlalchemy import text
|
||||
|
||||
from ..db import (
|
||||
cleanup_tables as _cleanup_tables,
|
||||
create_database,
|
||||
create_db_context,
|
||||
)
|
||||
|
||||
|
||||
async def cleanup_tables(
|
||||
session: AsyncSession,
|
||||
base: type[DeclarativeBase],
|
||||
) -> None:
|
||||
"""Truncate all tables for fast between-test cleanup.
|
||||
|
||||
.. deprecated::
|
||||
Import ``cleanup_tables`` from ``fastapi_toolsets.db`` instead.
|
||||
This re-export will be removed in v3.0.0.
|
||||
"""
|
||||
warnings.warn(
|
||||
"Importing cleanup_tables from fastapi_toolsets.pytest is deprecated "
|
||||
"and will be removed in v3.0.0. "
|
||||
"Use 'from fastapi_toolsets.db import cleanup_tables' instead.",
|
||||
DeprecationWarning,
|
||||
stacklevel=2,
|
||||
)
|
||||
await _cleanup_tables(session=session, base=base)
|
||||
|
||||
|
||||
def _get_xdist_worker(default_test_db: str) -> str:
|
||||
"""Return the pytest-xdist worker name, or *default_test_db* when not running under xdist.
|
||||
|
||||
Reads the ``PYTEST_XDIST_WORKER`` environment variable that xdist sets
|
||||
automatically in each worker process (e.g. ``"gw0"``, ``"gw1"``).
|
||||
When xdist is not installed or not active, the variable is absent and
|
||||
*default_test_db* is returned instead.
|
||||
|
||||
Args:
|
||||
default_test_db: Fallback value returned when ``PYTEST_XDIST_WORKER``
|
||||
is not set.
|
||||
"""
|
||||
return os.environ.get("PYTEST_XDIST_WORKER", default_test_db)
|
||||
|
||||
|
||||
def worker_database_url(database_url: str, default_test_db: str) -> str:
|
||||
"""Derive a per-worker database URL for pytest-xdist parallel runs.
|
||||
|
||||
Appends ``_{worker_name}`` to the database name so each xdist worker
|
||||
operates on its own database. When not running under xdist,
|
||||
``_{default_test_db}`` is appended instead.
|
||||
|
||||
The worker name is read from the ``PYTEST_XDIST_WORKER`` environment
|
||||
variable (set automatically by xdist in each worker process).
|
||||
|
||||
Args:
|
||||
database_url: Original database connection URL.
|
||||
default_test_db: Suffix appended to the database name when
|
||||
``PYTEST_XDIST_WORKER`` is not set.
|
||||
|
||||
Returns:
|
||||
A database URL with a worker- or default-specific database name.
|
||||
"""
|
||||
worker = _get_xdist_worker(default_test_db=default_test_db)
|
||||
|
||||
url = make_url(database_url)
|
||||
url = url.set(database=f"{url.database}_{worker}")
|
||||
return url.render_as_string(hide_password=False)
|
||||
|
||||
|
||||
@asynccontextmanager
|
||||
async def create_worker_database(
|
||||
database_url: str,
|
||||
default_test_db: str = "test_db",
|
||||
) -> AsyncGenerator[str, None]:
|
||||
"""Create and drop a per-worker database for pytest-xdist isolation.
|
||||
|
||||
Derives a worker-specific database URL using :func:`worker_database_url`,
|
||||
then delegates to :func:`~fastapi_toolsets.db.create_database` to create
|
||||
and drop it. Intended for use as a **session-scoped** fixture.
|
||||
|
||||
When running under xdist the database name is suffixed with the worker
|
||||
name (e.g. ``_gw0``). Otherwise it is suffixed with *default_test_db*.
|
||||
|
||||
Args:
|
||||
database_url: Original database connection URL (used as the server
|
||||
connection and as the base for the worker database name).
|
||||
default_test_db: Suffix appended to the database name when
|
||||
``PYTEST_XDIST_WORKER`` is not set. Defaults to ``"test_db"``.
|
||||
|
||||
Yields:
|
||||
The worker-specific database URL.
|
||||
|
||||
Example:
|
||||
```python
|
||||
from fastapi_toolsets.pytest import create_worker_database, create_db_session
|
||||
|
||||
DATABASE_URL = "postgresql+asyncpg://postgres:postgres@localhost/test_db"
|
||||
|
||||
@pytest.fixture(scope="session")
|
||||
async def worker_db_url():
|
||||
async with create_worker_database(DATABASE_URL) as url:
|
||||
yield url
|
||||
|
||||
@pytest.fixture
|
||||
async def db_session(worker_db_url):
|
||||
async with create_db_session(
|
||||
worker_db_url, Base, cleanup=True
|
||||
) as session:
|
||||
yield session
|
||||
```
|
||||
"""
|
||||
worker_url = worker_database_url(
|
||||
database_url=database_url, default_test_db=default_test_db
|
||||
)
|
||||
worker_db_name: str = make_url(worker_url).database # type: ignore[assignment]
|
||||
|
||||
engine = create_async_engine(database_url, isolation_level="AUTOCOMMIT")
|
||||
try:
|
||||
async with engine.connect() as conn:
|
||||
await conn.execute(text(f"DROP DATABASE IF EXISTS {worker_db_name}"))
|
||||
await create_database(db_name=worker_db_name, server_url=database_url)
|
||||
|
||||
yield worker_url
|
||||
|
||||
async with engine.connect() as conn:
|
||||
await conn.execute(text(f"DROP DATABASE IF EXISTS {worker_db_name}"))
|
||||
finally:
|
||||
await engine.dispose()
|
||||
|
||||
|
||||
@asynccontextmanager
|
||||
@@ -156,160 +283,3 @@ async def create_db_session(
|
||||
await conn.run_sync(base.metadata.drop_all)
|
||||
finally:
|
||||
await engine.dispose()
|
||||
|
||||
|
||||
def _get_xdist_worker(default_test_db: str) -> str:
|
||||
"""Return the pytest-xdist worker name, or *default_test_db* when not running under xdist.
|
||||
|
||||
Reads the ``PYTEST_XDIST_WORKER`` environment variable that xdist sets
|
||||
automatically in each worker process (e.g. ``"gw0"``, ``"gw1"``).
|
||||
When xdist is not installed or not active, the variable is absent and
|
||||
*default_test_db* is returned instead.
|
||||
|
||||
Args:
|
||||
default_test_db: Fallback value returned when ``PYTEST_XDIST_WORKER``
|
||||
is not set.
|
||||
"""
|
||||
return os.environ.get("PYTEST_XDIST_WORKER", default_test_db)
|
||||
|
||||
|
||||
def worker_database_url(database_url: str, default_test_db: str) -> str:
|
||||
"""Derive a per-worker database URL for pytest-xdist parallel runs.
|
||||
|
||||
Appends ``_{worker_name}`` to the database name so each xdist worker
|
||||
operates on its own database. When not running under xdist,
|
||||
``_{default_test_db}`` is appended instead.
|
||||
|
||||
The worker name is read from the ``PYTEST_XDIST_WORKER`` environment
|
||||
variable (set automatically by xdist in each worker process).
|
||||
|
||||
Args:
|
||||
database_url: Original database connection URL.
|
||||
default_test_db: Suffix appended to the database name when
|
||||
``PYTEST_XDIST_WORKER`` is not set.
|
||||
|
||||
Returns:
|
||||
A database URL with a worker- or default-specific database name.
|
||||
|
||||
Example:
|
||||
```python
|
||||
# With PYTEST_XDIST_WORKER="gw0":
|
||||
url = worker_database_url(
|
||||
"postgresql+asyncpg://user:pass@localhost/test_db",
|
||||
default_test_db="test",
|
||||
)
|
||||
# "postgresql+asyncpg://user:pass@localhost/test_db_gw0"
|
||||
|
||||
# Without PYTEST_XDIST_WORKER:
|
||||
url = worker_database_url(
|
||||
"postgresql+asyncpg://user:pass@localhost/test_db",
|
||||
default_test_db="test",
|
||||
)
|
||||
# "postgresql+asyncpg://user:pass@localhost/test_db_test"
|
||||
```
|
||||
"""
|
||||
worker = _get_xdist_worker(default_test_db=default_test_db)
|
||||
|
||||
url = make_url(database_url)
|
||||
url = url.set(database=f"{url.database}_{worker}")
|
||||
return url.render_as_string(hide_password=False)
|
||||
|
||||
|
||||
@asynccontextmanager
|
||||
async def create_worker_database(
|
||||
database_url: str,
|
||||
default_test_db: str = "test_db",
|
||||
) -> AsyncGenerator[str, None]:
|
||||
"""Create and drop a per-worker database for pytest-xdist isolation.
|
||||
|
||||
Intended for use as a **session-scoped** fixture. Connects to the server
|
||||
using the original *database_url* (with ``AUTOCOMMIT`` isolation for DDL),
|
||||
creates a dedicated database for the worker, and yields the worker-specific
|
||||
URL. On cleanup the worker database is dropped.
|
||||
|
||||
When running under xdist the database name is suffixed with the worker
|
||||
name (e.g. ``_gw0``). Otherwise it is suffixed with *default_test_db*.
|
||||
|
||||
Args:
|
||||
database_url: Original database connection URL.
|
||||
default_test_db: Suffix appended to the database name when
|
||||
``PYTEST_XDIST_WORKER`` is not set. Defaults to ``"test_db"``.
|
||||
|
||||
Yields:
|
||||
The worker-specific database URL.
|
||||
|
||||
Example:
|
||||
```python
|
||||
from fastapi_toolsets.pytest import (
|
||||
create_worker_database, create_db_session,
|
||||
)
|
||||
|
||||
DATABASE_URL = "postgresql+asyncpg://postgres:postgres@localhost/test_db"
|
||||
|
||||
@pytest.fixture(scope="session")
|
||||
async def worker_db_url():
|
||||
async with create_worker_database(DATABASE_URL) as url:
|
||||
yield url
|
||||
|
||||
@pytest.fixture
|
||||
async def db_session(worker_db_url):
|
||||
async with create_db_session(
|
||||
worker_db_url, Base, cleanup=True
|
||||
) as session:
|
||||
yield session
|
||||
```
|
||||
"""
|
||||
worker_url = worker_database_url(
|
||||
database_url=database_url, default_test_db=default_test_db
|
||||
)
|
||||
worker_db_name = make_url(worker_url).database
|
||||
|
||||
engine = create_async_engine(
|
||||
database_url,
|
||||
isolation_level="AUTOCOMMIT",
|
||||
)
|
||||
try:
|
||||
async with engine.connect() as conn:
|
||||
await conn.execute(text(f"DROP DATABASE IF EXISTS {worker_db_name}"))
|
||||
await conn.execute(text(f"CREATE DATABASE {worker_db_name}"))
|
||||
|
||||
yield worker_url
|
||||
|
||||
async with engine.connect() as conn:
|
||||
await conn.execute(text(f"DROP DATABASE IF EXISTS {worker_db_name}"))
|
||||
finally:
|
||||
await engine.dispose()
|
||||
|
||||
|
||||
async def cleanup_tables(
|
||||
session: AsyncSession,
|
||||
base: type[DeclarativeBase],
|
||||
) -> None:
|
||||
"""Truncate all tables for fast between-test cleanup.
|
||||
|
||||
Executes a single ``TRUNCATE … RESTART IDENTITY CASCADE`` statement
|
||||
across every table in *base*'s metadata, which is significantly faster
|
||||
than dropping and re-creating tables between tests.
|
||||
|
||||
This is a no-op when the metadata contains no tables.
|
||||
|
||||
Args:
|
||||
session: An active async database session.
|
||||
base: SQLAlchemy DeclarativeBase class containing model metadata.
|
||||
|
||||
Example:
|
||||
```python
|
||||
@pytest.fixture
|
||||
async def db_session(worker_db_url):
|
||||
async with create_db_session(worker_db_url, Base) as session:
|
||||
yield session
|
||||
await cleanup_tables(session, Base)
|
||||
```
|
||||
"""
|
||||
tables = base.metadata.sorted_tables
|
||||
if not tables:
|
||||
return
|
||||
|
||||
table_names = ", ".join(f'"{t.name}"' for t in tables)
|
||||
await session.execute(text(f"TRUNCATE {table_names} RESTART IDENTITY CASCADE"))
|
||||
await session.commit()
|
||||
|
||||
Reference in New Issue
Block a user