Files
fastapi-toolsets/tests/test_db.py
T

1500 lines
56 KiB
Python

"""Tests for fastapi_toolsets.db module (v5 ``Database`` facade)."""
import asyncio
import uuid
from contextlib import asynccontextmanager
from unittest.mock import AsyncMock, patch
import pytest
from fastapi import Depends, FastAPI
from fastapi.responses import StreamingResponse
from httpx import ASGITransport, AsyncClient
from sqlalchemy import (
Column,
ForeignKey,
ForeignKeyConstraint,
String,
Table,
Uuid,
select,
text,
)
from sqlalchemy.engine import make_url
from sqlalchemy.exc import IntegrityError
from sqlalchemy.ext.asyncio import (
AsyncEngine,
AsyncSession,
async_sessionmaker,
create_async_engine,
)
from sqlalchemy.orm import (
DeclarativeBase,
Mapped,
mapped_column,
relationship,
selectinload,
)
from starlette.requests import Request
from fastapi_toolsets.db import (
Database,
LockMode,
advisory_lock,
lock_tables,
m2m_add,
m2m_remove,
m2m_set,
transaction,
wait_for_row_change,
)
from fastapi_toolsets.db.core import _CommitOnResponseMiddleware
from fastapi_toolsets.db.testing import cleanup_tables, create_database
from fastapi_toolsets.exceptions import (
LockTimeoutError,
NotFoundError,
PoolExhaustedError,
)
from fastapi_toolsets.pytest import create_db_session
from .conftest import (
DATABASE_URL,
Base,
Post,
Role,
RoleCreate,
RoleCrud,
Tag,
User,
UserCrud,
)
def _make_request() -> Request:
"""Minimal ASGI HTTP request for exercising the Database dependency directly."""
return Request({"type": "http", "headers": []})
class TestDatabaseConstruction:
"""Construction contract: provide exactly one of url / engine."""
def test_requires_url_or_engine(self):
"""Neither url nor engine raises TypeError."""
with pytest.raises(TypeError):
Database()
@pytest.mark.anyio
async def test_both_url_and_engine_raises(self, engine):
"""Both url and engine raises TypeError."""
with pytest.raises(TypeError):
Database(DATABASE_URL, engine=engine)
@pytest.mark.anyio
async def test_engine_options_with_engine_raises(self, engine):
"""engine_options are rejected in engine= mode."""
with pytest.raises(TypeError):
Database(engine=engine, pool_size=5)
@pytest.mark.anyio
async def test_url_mode_owns_engine(self):
"""URL mode builds and owns the engine."""
db = Database(DATABASE_URL)
try:
assert db._owns_engine is True
assert db.engine is not None
finally:
await db.engine.dispose()
@pytest.mark.anyio
async def test_engine_mode_borrows_engine(self, engine):
"""engine= mode reuses the given engine and does not own it."""
db = Database(engine=engine)
assert db._owns_engine is False
assert db.engine is engine
@pytest.mark.anyio
async def test_distinct_instances_use_distinct_state_attrs(self, engine):
"""Two Database instances never share a request-state attribute."""
a = Database(engine=engine)
b = Database(engine=engine)
assert a._state_attr != b._state_attr
@pytest.mark.anyio
async def test_lifespan_disposes_owned_engine(self):
"""The lifespan disposes the engine it built (URL mode)."""
db = Database(DATABASE_URL)
# ``AsyncEngine.dispose`` is read-only on the instance, so patch the class.
with patch.object(AsyncEngine, "dispose", new=AsyncMock()) as disposed:
async with db.lifespan(None):
pass
disposed.assert_awaited_once()
await db.engine.dispose()
@pytest.mark.anyio
async def test_lifespan_skips_borrowed_engine(self):
"""The lifespan leaves a borrowed engine untouched (engine= mode)."""
eng = create_async_engine(DATABASE_URL, echo=False)
db = Database(engine=eng)
with patch.object(AsyncEngine, "dispose", new=AsyncMock()) as disposed:
async with db.lifespan(None):
pass
disposed.assert_not_awaited()
await eng.dispose()
class TestLifespanComposition:
"""``install`` composes engine disposal around the app's own lifespan."""
@pytest.mark.anyio
async def test_install_composes_user_lifespan(self):
"""A user-defined lifespan runs, and the engine is disposed after it."""
events: list[str] = []
@asynccontextmanager
async def user_lifespan(app):
events.append("startup")
yield
events.append("shutdown")
db = Database(DATABASE_URL)
app = FastAPI(lifespan=user_lifespan)
db.install(app)
with patch.object(AsyncEngine, "dispose", new=AsyncMock()) as disposed:
async with app.router.lifespan_context(app):
assert events == ["startup"]
disposed.assert_not_awaited()
# User shutdown runs, then the engine is disposed.
assert events == ["startup", "shutdown"]
disposed.assert_awaited_once()
await db.engine.dispose()
@pytest.mark.anyio
async def test_install_disposes_without_user_lifespan(self):
"""``install`` disposes the engine even when the app has no custom lifespan."""
db = Database(DATABASE_URL)
app = FastAPI()
db.install(app)
with patch.object(AsyncEngine, "dispose", new=AsyncMock()) as disposed:
async with app.router.lifespan_context(app):
disposed.assert_not_awaited()
disposed.assert_awaited_once()
await db.engine.dispose()
@pytest.mark.anyio
async def test_disposal_is_idempotent(self):
"""Combining ``lifespan=db.lifespan`` with ``install`` disposes only once."""
db = Database(DATABASE_URL)
app = FastAPI(lifespan=db.lifespan)
db.install(app)
with patch.object(AsyncEngine, "dispose", new=AsyncMock()) as disposed:
async with app.router.lifespan_context(app):
pass
disposed.assert_awaited_once()
await db.engine.dispose()
@pytest.mark.anyio
async def test_install_skips_disposal_for_borrowed_engine(self, engine):
"""``install`` never disposes an engine it does not own."""
db = Database(engine=engine)
app = FastAPI()
db.install(app)
with patch.object(AsyncEngine, "dispose", new=AsyncMock()) as disposed:
async with app.router.lifespan_context(app):
pass
disposed.assert_not_awaited()
class TestDatabaseDependency:
"""Tests for the FastAPI dependency (``Depends(db)`` / ``db.__call__``)."""
@pytest.mark.anyio
async def test_yields_session(self, engine):
"""Dependency yields a valid session."""
db = Database(engine=engine)
async for session in db(_make_request()):
assert isinstance(session, AsyncSession)
break
@pytest.mark.anyio
async def test_auto_commits_transaction(self, engine, session_maker):
"""Without middleware, the dependency commits an open transaction on exit."""
db = Database(engine=engine)
async for session in db(_make_request()):
role = Role(name="test_role_dep")
session.add(role)
await session.flush()
async with session_maker() as verify:
result = await RoleCrud.first(verify, [Role.name == "test_role_dep"])
assert result is not None
@pytest.mark.anyio
async def test_in_transaction_on_yield(self, engine):
"""Session is already in a transaction when the endpoint body starts."""
db = Database(engine=engine)
async for session in db(_make_request()):
assert session.in_transaction()
break
@pytest.mark.anyio
async def test_no_commit_when_not_in_transaction(self, engine):
"""Dependency skips commit if the session left its transaction on exit."""
db = Database(engine=engine)
async for session in db(_make_request()):
await session.commit()
assert not session.in_transaction()
# The post-yield path must not call commit again (no error).
@pytest.mark.anyio
async def test_stashes_session_on_request_state(self, engine):
"""Dependency exposes the session on request.state for the commit middleware."""
db = Database(engine=engine)
request = _make_request()
async for session in db(request):
assert getattr(request.state, db._state_attr) is session
break
@pytest.mark.anyio
async def test_skips_commit_when_middleware_installed(self, engine, session_maker):
"""With ``install()``, the dependency must NOT commit — the middleware owns it.
Here no middleware actually runs (we call the dependency directly), so the
open transaction is rolled back on session close and nothing persists.
"""
db = Database(engine=engine)
db.install(FastAPI())
async for session in db(_make_request()):
role = Role(name="mw_owns_commit")
session.add(role)
await session.flush()
async with session_maker() as verify:
result = await RoleCrud.first(verify, [Role.name == "mw_owns_commit"])
assert result is None
class TestDatabaseSession:
"""Tests for ``db.session()`` (sessions outside request handlers)."""
@pytest.mark.anyio
async def test_context_manager_yields_session(self, engine):
"""Context manager yields a valid session."""
db = Database(engine=engine)
async with db.session() as session:
assert isinstance(session, AsyncSession)
@pytest.mark.anyio
async def test_context_manager_commits(self, engine, session_maker):
"""Context manager commits on exit."""
db = Database(engine=engine)
async with db.session() as session:
role = Role(name="context_role")
session.add(role)
await session.flush()
async with session_maker() as verify:
result = await RoleCrud.first(verify, [Role.name == "context_role"])
assert result is not None
@pytest.mark.anyio
async def test_no_commit_when_not_in_transaction(self, engine):
"""Context skips commit if the session left its transaction on exit."""
db = Database(engine=engine)
async with db.session() as session:
await session.commit()
assert not session.in_transaction()
@pytest.mark.anyio
async def test_pool_exhausted_raises_pool_exhausted_error(self):
"""PoolExhaustedError is raised when the pool is exhausted on session entry."""
db = Database(DATABASE_URL, pool_size=1, max_overflow=0, pool_timeout=0.1)
try:
async with db.session(): # checks out the single available connection
with pytest.raises(PoolExhaustedError):
async with db.session():
pass
finally:
await db.engine.dispose()
class TestDatabaseBegin:
"""Tests for ``db.begin()`` (open a session already in a transaction)."""
@pytest.mark.anyio
async def test_commits_on_success(self, engine, session_maker):
"""The block commits when it exits cleanly."""
db = Database(engine=engine)
async with db.begin() as session:
session.add(Role(name="begin_role"))
async with session_maker() as verify:
result = await RoleCrud.first(verify, [Role.name == "begin_role"])
assert result is not None
@pytest.mark.anyio
async def test_rolls_back_on_exception(self, engine, session_maker):
"""The block rolls back on exception."""
db = Database(engine=engine)
with pytest.raises(ValueError):
async with db.begin() as session:
session.add(Role(name="begin_rollback_role"))
await session.flush()
raise ValueError("Simulated error")
async with session_maker() as verify:
result = await RoleCrud.first(verify, [Role.name == "begin_rollback_role"])
assert result is None
class TestTransaction:
"""Tests for the ``transaction`` context manager (savepoint-aware primitive)."""
@pytest.mark.anyio
async def test_starts_transaction(self, db_session: AsyncSession):
"""transaction starts a new transaction."""
async with transaction(db_session):
role = Role(name="tx_role")
db_session.add(role)
result = await RoleCrud.first(db_session, [Role.name == "tx_role"])
assert result is not None
@pytest.mark.anyio
async def test_nested_transaction_uses_savepoint(self, db_session: AsyncSession):
"""Nested transactions use savepoints."""
async with transaction(db_session):
role1 = Role(name="outer_role")
db_session.add(role1)
await db_session.flush()
async with transaction(db_session):
role2 = Role(name="inner_role")
db_session.add(role2)
results = await RoleCrud.get_multi(db_session)
names = {r.name for r in results}
assert "outer_role" in names
assert "inner_role" in names
@pytest.mark.anyio
async def test_rollback_on_exception(self, db_session: AsyncSession):
"""Transaction rolls back on exception."""
try:
async with transaction(db_session):
role = Role(name="rollback_role")
db_session.add(role)
await db_session.flush()
raise ValueError("Simulated error")
except ValueError:
pass
result = await RoleCrud.first(db_session, [Role.name == "rollback_role"])
assert result is None
@pytest.mark.anyio
async def test_nested_rollback_preserves_outer(self, db_session: AsyncSession):
"""Nested rollback preserves outer transaction."""
async with transaction(db_session):
role1 = Role(name="preserved_role")
db_session.add(role1)
await db_session.flush()
try:
async with transaction(db_session):
role2 = Role(name="rolled_back_role")
db_session.add(role2)
await db_session.flush()
raise ValueError("Inner error")
except ValueError:
pass
outer = await RoleCrud.first(db_session, [Role.name == "preserved_role"])
inner = await RoleCrud.first(db_session, [Role.name == "rolled_back_role"])
assert outer is not None
assert inner is None
class TestLockMode:
"""Tests for LockMode enum."""
def test_lock_modes_exist(self):
"""All expected lock modes are defined."""
assert LockMode.ACCESS_SHARE == "ACCESS SHARE"
assert LockMode.ROW_SHARE == "ROW SHARE"
assert LockMode.ROW_EXCLUSIVE == "ROW EXCLUSIVE"
assert LockMode.SHARE_UPDATE_EXCLUSIVE == "SHARE UPDATE EXCLUSIVE"
assert LockMode.SHARE == "SHARE"
assert LockMode.SHARE_ROW_EXCLUSIVE == "SHARE ROW EXCLUSIVE"
assert LockMode.EXCLUSIVE == "EXCLUSIVE"
assert LockMode.ACCESS_EXCLUSIVE == "ACCESS EXCLUSIVE"
def test_lock_mode_is_string(self):
"""Lock modes are string enums."""
assert isinstance(LockMode.EXCLUSIVE, str)
assert LockMode.EXCLUSIVE.value == "EXCLUSIVE"
class TestLockTables:
"""Tests for ``db.lock_tables`` (PostgreSQL-specific)."""
@pytest.mark.anyio
async def test_lock_single_table(self, engine, session_maker):
"""Lock a single table; changes inside are committed on context exit."""
db = Database(engine=engine)
async with db.lock_tables([Role]) as session:
role = Role(name="locked_role")
session.add(role)
async with session_maker() as verify:
result = await RoleCrud.first(verify, [Role.name == "locked_role"])
assert result is not None
@pytest.mark.anyio
async def test_lock_multiple_tables(self, engine, session_maker):
"""Lock multiple tables."""
db = Database(engine=engine)
async with db.lock_tables([Role, User]) as session:
role = Role(name="multi_lock_role")
session.add(role)
async with session_maker() as verify:
result = await RoleCrud.first(verify, [Role.name == "multi_lock_role"])
assert result is not None
@pytest.mark.anyio
async def test_lock_with_custom_mode(self, engine, session_maker):
"""Lock with custom lock mode."""
db = Database(engine=engine)
async with db.lock_tables([Role], mode=LockMode.EXCLUSIVE) as session:
role = Role(name="exclusive_lock_role")
session.add(role)
async with session_maker() as verify:
result = await RoleCrud.first(verify, [Role.name == "exclusive_lock_role"])
assert result is not None
@pytest.mark.anyio
async def test_lock_rollback_on_exception(self, engine, session_maker):
"""Lock context rolls back on exception."""
db = Database(engine=engine)
with pytest.raises(ValueError):
async with db.lock_tables([Role]) as session:
role = Role(name="lock_rollback_role")
session.add(role)
await session.flush()
raise ValueError("Simulated error")
async with session_maker() as verify:
result = await RoleCrud.first(verify, [Role.name == "lock_rollback_role"])
assert result is None
class TestAdvisoryLock:
"""Tests for advisory_lock context manager (PostgreSQL-specific)."""
@pytest.mark.anyio
async def test_blocking_exclusive_acquires(self, db_session: AsyncSession):
"""Blocking exclusive lock acquires and yields True."""
async with advisory_lock(db_session, 1001) as acquired:
assert acquired is True
@pytest.mark.anyio
async def test_nowait_returns_true_when_free(self, db_session: AsyncSession):
"""nowait=True yields True when the lock is available."""
async with advisory_lock(db_session, 1002, nowait=True) as acquired:
assert acquired is True
@pytest.mark.anyio
async def test_nowait_returns_false_when_contended(self, session_maker):
"""nowait=True yields False when another session holds the lock."""
async with session_maker() as holder:
async with holder.begin():
async with advisory_lock(holder, 1003):
async with session_maker() as contender:
async with contender.begin():
async with advisory_lock(
contender, 1003, nowait=True
) as acquired:
assert acquired is False
@pytest.mark.anyio
async def test_shared_allows_concurrent_readers(self, session_maker):
"""Two shared locks on the same key are both acquired."""
async with session_maker() as s1, session_maker() as s2:
async with s1.begin(), s2.begin():
async with advisory_lock(s1, 1004, shared=True) as a1:
async with advisory_lock(s2, 1004, shared=True, nowait=True) as a2:
assert a1 is True
assert a2 is True
@pytest.mark.anyio
async def test_acquire_does_not_flush_pending(self, db_session: AsyncSession):
"""Acquiring the lock must not autoflush the caller's pending ORM changes.
Guards the SQLAlchemy 2.1 behavior where raw ``text()`` autoflushes too;
the helper wraps lock SQL in ``no_autoflush`` to preserve v4 semantics.
"""
role = Role(name="not_flushed_by_lock")
db_session.add(role)
async with advisory_lock(db_session, 2001):
# The pending INSERT must still be unflushed inside the lock.
assert role in db_session.new
@pytest.mark.anyio
async def test_tuple_key(self, db_session: AsyncSession):
"""(int, int) key variant acquires the lock."""
async with advisory_lock(db_session, (7, 42)) as acquired:
assert acquired is True
@pytest.mark.anyio
async def test_tuple_key_nowait_contended(self, session_maker):
"""Tuple key nowait returns False when contended."""
async with session_maker() as holder:
async with holder.begin():
async with advisory_lock(holder, (7, 99)):
async with session_maker() as contender:
async with contender.begin():
async with advisory_lock(
contender, (7, 99), nowait=True
) as acquired:
assert acquired is False
@pytest.mark.anyio
async def test_lock_released_at_context_exit(self, session_maker):
"""Lock is released when the context exits, even while the transaction is still open."""
async with session_maker() as s1:
async with s1.begin():
async with advisory_lock(s1, 1005):
pass # lock released here — transaction still active
async with session_maker() as s2:
async with s2.begin():
async with advisory_lock(s2, 1005, nowait=True) as acquired:
assert (
acquired is True
) # s1 still in transaction but lock is free
@pytest.mark.anyio
async def test_timeout_raises_when_contended(self, session_maker):
"""timeout= raises LockTimeoutError when the lock cannot be acquired."""
async with session_maker() as holder:
async with holder.begin():
async with advisory_lock(holder, 1006):
async with session_maker() as contender:
async with contender.begin():
with pytest.raises(LockTimeoutError):
async with advisory_lock(
contender, 1006, timeout="10ms"
):
pass
class TestWaitForRowChange:
"""Tests for wait_for_row_change polling function."""
@pytest.mark.anyio
async def test_detects_update(self, db_session: AsyncSession, engine):
"""Returns updated instance when a column value changes."""
role = Role(name="watch_role")
db_session.add(role)
await db_session.commit()
async def update_later():
await asyncio.sleep(0.15)
factory = async_sessionmaker(engine, expire_on_commit=False)
async with factory() as other:
r = await other.get(Role, role.id)
assert r is not None
r.name = "updated_role"
await other.commit()
update_task = asyncio.create_task(update_later())
result = await wait_for_row_change(db_session, Role, role.id, interval=0.05)
await update_task
assert result.name == "updated_role"
@pytest.mark.anyio
async def test_watches_specific_columns(self, db_session: AsyncSession, engine):
"""Only triggers on changes to specified columns."""
user = User(username="testuser", email="test@example.com")
db_session.add(user)
await db_session.commit()
async def update_later():
factory = async_sessionmaker(engine, expire_on_commit=False)
# First: change email (not watched) — should not trigger
await asyncio.sleep(0.15)
async with factory() as other:
u = await other.get(User, user.id)
assert u is not None
u.email = "new@example.com"
await other.commit()
# Second: change username (watched) — should trigger
await asyncio.sleep(0.15)
async with factory() as other:
u = await other.get(User, user.id)
assert u is not None
u.username = "newuser"
await other.commit()
update_task = asyncio.create_task(update_later())
result = await wait_for_row_change(
db_session, User, user.id, columns=["username"], interval=0.05
)
await update_task
assert result.username == "newuser"
assert result.email == "new@example.com"
@pytest.mark.anyio
async def test_nonexistent_row_raises(self, db_session: AsyncSession):
"""Raises NotFoundError when the row does not exist."""
fake_id = uuid.uuid4()
with pytest.raises(NotFoundError, match="not found"):
await wait_for_row_change(db_session, Role, fake_id, interval=0.05)
@pytest.mark.anyio
async def test_timeout_raises(self, db_session: AsyncSession):
"""Raises TimeoutError when no change is detected within timeout."""
role = Role(name="timeout_role")
db_session.add(role)
await db_session.commit()
with pytest.raises(TimeoutError):
await wait_for_row_change(
db_session, Role, role.id, interval=0.05, timeout=0.2
)
@pytest.mark.anyio
async def test_deleted_row_raises(self, db_session: AsyncSession, engine):
"""Raises NotFoundError when the row is deleted during polling."""
role = Role(name="delete_role")
db_session.add(role)
await db_session.commit()
async def delete_later():
await asyncio.sleep(0.15)
factory = async_sessionmaker(engine, expire_on_commit=False)
async with factory() as other:
r = await other.get(Role, role.id)
await other.delete(r)
await other.commit()
delete_task = asyncio.create_task(delete_later())
with pytest.raises(NotFoundError):
await wait_for_row_change(db_session, Role, role.id, interval=0.05)
await delete_task
class TestCreateDatabase:
"""Tests for create_database."""
@pytest.mark.anyio
async def test_creates_database(self):
"""Database is created by create_database."""
target_url = (
make_url(DATABASE_URL)
.set(database="test_create_db_general")
.render_as_string(hide_password=False)
)
expected_db = make_url(target_url).database
assert expected_db is not None
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 {expected_db}"))
await create_database(db_name=expected_db, server_url=DATABASE_URL)
async with engine.connect() as conn:
result = await conn.execute(
text("SELECT 1 FROM pg_database WHERE datname = :name"),
{"name": expected_db},
)
assert result.scalar() == 1
# Cleanup
async with engine.connect() as conn:
await conn.execute(text(f"DROP DATABASE IF EXISTS {expected_db}"))
finally:
await engine.dispose()
class TestCleanupTables:
"""Tests for cleanup_tables helper."""
@pytest.mark.anyio
async def test_truncates_all_tables(self):
"""All table rows are removed after cleanup_tables."""
async with create_db_session(DATABASE_URL, Base, drop_tables=True) as session:
role = Role(id=uuid.uuid4(), name="cleanup_role")
session.add(role)
await session.flush()
user = User(
id=uuid.uuid4(),
username="cleanup_user",
email="cleanup@test.com",
role_id=role.id,
)
session.add(user)
await session.commit()
# Verify rows exist
roles_count = await RoleCrud.count(session)
users_count = await UserCrud.count(session)
assert roles_count == 1
assert users_count == 1
await cleanup_tables(session, Base)
# Verify tables are empty
roles_count = await RoleCrud.count(session)
users_count = await UserCrud.count(session)
assert roles_count == 0
assert users_count == 0
@pytest.mark.anyio
async def test_noop_for_empty_metadata(self):
"""cleanup_tables does not raise when metadata has no tables."""
class EmptyBase(DeclarativeBase):
pass
async with create_db_session(DATABASE_URL, Base, drop_tables=True) as session:
# Should not raise
await cleanup_tables(session, EmptyBase)
class TestM2MAdd:
"""Tests for m2m_add helper."""
@pytest.mark.anyio
async def test_adds_single_related(self, db_session: AsyncSession):
"""Associates one related instance via the secondary table."""
user = User(username="m2m_author", email="m2m@test.com")
db_session.add(user)
await db_session.flush()
post = Post(title="Post A", author_id=user.id)
tag = Tag(name="python")
db_session.add_all([post, tag])
await db_session.flush()
async with transaction(db_session):
await m2m_add(db_session, post, Post.tags, tag)
result = await db_session.execute(
select(Post).where(Post.id == post.id).options(selectinload(Post.tags))
)
loaded = result.scalar_one()
assert len(loaded.tags) == 1
assert loaded.tags[0].id == tag.id
@pytest.mark.anyio
async def test_adds_multiple_related(self, db_session: AsyncSession):
"""Associates multiple related instances in a single call."""
user = User(username="m2m_author2", email="m2m2@test.com")
db_session.add(user)
await db_session.flush()
post = Post(title="Post B", author_id=user.id)
tag1 = Tag(name="web")
tag2 = Tag(name="api")
tag3 = Tag(name="async")
db_session.add_all([post, tag1, tag2, tag3])
await db_session.flush()
async with transaction(db_session):
await m2m_add(db_session, post, Post.tags, tag1, tag2, tag3)
result = await db_session.execute(
select(Post).where(Post.id == post.id).options(selectinload(Post.tags))
)
loaded = result.scalar_one()
assert {t.id for t in loaded.tags} == {tag1.id, tag2.id, tag3.id}
@pytest.mark.anyio
async def test_noop_for_empty_related(self, db_session: AsyncSession):
"""Calling with no related instances is a no-op."""
user = User(username="m2m_author3", email="m2m3@test.com")
db_session.add(user)
await db_session.flush()
post = Post(title="Post C", author_id=user.id)
db_session.add(post)
await db_session.flush()
async with transaction(db_session):
await m2m_add(db_session, post, Post.tags) # no related instances
result = await db_session.execute(
select(Post).where(Post.id == post.id).options(selectinload(Post.tags))
)
loaded = result.scalar_one()
assert loaded.tags == []
@pytest.mark.anyio
async def test_ignore_conflicts_true(self, db_session: AsyncSession):
"""Duplicate inserts are silently skipped when ignore_conflicts=True."""
user = User(username="m2m_author4", email="m2m4@test.com")
db_session.add(user)
await db_session.flush()
post = Post(title="Post D", author_id=user.id)
tag = Tag(name="duplicate_tag")
db_session.add_all([post, tag])
await db_session.flush()
async with transaction(db_session):
await m2m_add(db_session, post, Post.tags, tag)
# Second call with ignore_conflicts=True must not raise
async with transaction(db_session):
await m2m_add(db_session, post, Post.tags, tag, ignore_conflicts=True)
result = await db_session.execute(
select(Post).where(Post.id == post.id).options(selectinload(Post.tags))
)
loaded = result.scalar_one()
assert len(loaded.tags) == 1
@pytest.mark.anyio
async def test_ignore_conflicts_false_raises(self, db_session: AsyncSession):
"""Duplicate inserts raise IntegrityError when ignore_conflicts=False (default)."""
user = User(username="m2m_author5", email="m2m5@test.com")
db_session.add(user)
await db_session.flush()
post = Post(title="Post E", author_id=user.id)
tag = Tag(name="conflict_tag")
db_session.add_all([post, tag])
await db_session.flush()
async with transaction(db_session):
await m2m_add(db_session, post, Post.tags, tag)
with pytest.raises(IntegrityError):
async with transaction(db_session):
await m2m_add(db_session, post, Post.tags, tag)
@pytest.mark.anyio
async def test_non_m2m_raises_type_error(self, db_session: AsyncSession):
"""Passing a non-M2M relationship attribute raises TypeError."""
user = User(username="m2m_author6", email="m2m6@test.com")
db_session.add(user)
await db_session.flush()
role = Role(name="type_err_role")
db_session.add(role)
await db_session.flush()
with pytest.raises(TypeError, match="Many-to-Many"):
await m2m_add(db_session, user, User.role, role)
@pytest.mark.anyio
async def test_works_inside_lock_tables(self, session_maker):
"""m2m_add works correctly inside a lock_tables context."""
async with lock_tables(session_maker, [Tag]) as session:
user = User(username="m2m_lock_author", email="m2m_lock@test.com")
session.add(user)
await session.flush()
tag = Tag(name="locked_tag")
session.add(tag)
await session.flush()
post = Post(title="Post Lock", author_id=user.id)
session.add(post)
await session.flush()
await m2m_add(session, post, Post.tags, tag)
async with session_maker() as verify:
result = await verify.execute(
select(Post).where(Post.id == post.id).options(selectinload(Post.tags))
)
loaded = result.scalar_one()
assert len(loaded.tags) == 1
assert loaded.tags[0].name == "locked_tag"
class TestDbErrors:
"""Tests for structured error handling in db utilities."""
@pytest.mark.anyio
async def test_pool_exhausted_on_dependency_raises_pool_exhausted_error(self):
"""PoolExhaustedError is raised when the pool is exhausted on the dependency."""
db = Database(DATABASE_URL, pool_size=1, max_overflow=0, pool_timeout=0.1)
try:
async with db.session(): # check out the single available connection
with pytest.raises(PoolExhaustedError):
async for _ in db(_make_request()):
pass
finally:
await db.engine.dispose()
@pytest.mark.anyio
async def test_pool_exhausted_on_lock_tables_raises_pool_exhausted_error(self):
"""PoolExhaustedError is raised when the pool is exhausted on lock_tables."""
db = Database(DATABASE_URL, pool_size=1, max_overflow=0, pool_timeout=0.1)
try:
async with db.session(): # check out the single available connection
with pytest.raises(PoolExhaustedError):
async with db.lock_tables([Role]):
pass
finally:
await db.engine.dispose()
@pytest.mark.anyio
async def test_lock_timeout_raises_lock_timeout_error(self, engine, session_maker):
"""LockTimeoutError is raised when a table lock cannot be acquired within timeout."""
db = Database(engine=engine)
async with db.lock_tables([Role]):
with pytest.raises(LockTimeoutError):
async with db.lock_tables([Role], timeout="100ms"):
pass
class _LocalBase(DeclarativeBase):
pass
_comp_assoc = Table(
"_comp_assoc",
_LocalBase.metadata,
Column("owner_id", Uuid, ForeignKey("_comp_owners.id"), primary_key=True),
Column("item_group", String(50), primary_key=True),
Column("item_code", String(50), primary_key=True),
ForeignKeyConstraint(
["item_group", "item_code"],
["_comp_items.group_id", "_comp_items.item_code"],
),
)
class _CompOwner(_LocalBase):
__tablename__ = "_comp_owners"
id: Mapped[uuid.UUID] = mapped_column(Uuid, primary_key=True, default=uuid.uuid4)
items: Mapped[list["_CompItem"]] = relationship(secondary=_comp_assoc)
class _CompItem(_LocalBase):
__tablename__ = "_comp_items"
group_id: Mapped[str] = mapped_column(String(50), primary_key=True)
item_code: Mapped[str] = mapped_column(String(50), primary_key=True)
class TestM2MRemove:
"""Tests for m2m_remove helper."""
async def _setup(
self, session: AsyncSession, username: str, email: str, *tag_names: str
):
"""Create a user, post, and tags; associate all tags with the post."""
user = User(username=username, email=email)
session.add(user)
await session.flush()
post = Post(title=f"Post {username}", author_id=user.id)
tags = [Tag(name=n) for n in tag_names]
session.add(post)
session.add_all(tags)
await session.flush()
async with transaction(session):
await m2m_add(session, post, Post.tags, *tags)
return post, tags
async def _load_tags(self, session: AsyncSession, post: Post) -> list[Tag]:
result = await session.execute(
select(Post).where(Post.id == post.id).options(selectinload(Post.tags))
)
return result.scalar_one().tags
@pytest.mark.anyio
async def test_removes_single(self, db_session: AsyncSession):
"""Removes one association, leaving others intact."""
post, (tag1, tag2) = await self._setup(
db_session, "rm_author1", "rm1@test.com", "tag_rm_a", "tag_rm_b"
)
async with transaction(db_session):
await m2m_remove(db_session, post, Post.tags, tag1)
remaining = await self._load_tags(db_session, post)
assert len(remaining) == 1
assert remaining[0].id == tag2.id
@pytest.mark.anyio
async def test_removes_multiple(self, db_session: AsyncSession):
"""Removes multiple associations in one call."""
post, (tag1, tag2, tag3) = await self._setup(
db_session, "rm_author2", "rm2@test.com", "tag_rm_c", "tag_rm_d", "tag_rm_e"
)
async with transaction(db_session):
await m2m_remove(db_session, post, Post.tags, tag1, tag3)
remaining = await self._load_tags(db_session, post)
assert len(remaining) == 1
assert remaining[0].id == tag2.id
@pytest.mark.anyio
async def test_noop_for_empty_related(self, db_session: AsyncSession):
"""Calling with no related instances is a no-op."""
post, (tag,) = await self._setup(
db_session, "rm_author3", "rm3@test.com", "tag_rm_f"
)
async with transaction(db_session):
await m2m_remove(db_session, post, Post.tags)
remaining = await self._load_tags(db_session, post)
assert len(remaining) == 1
@pytest.mark.anyio
async def test_idempotent_for_missing_association(self, db_session: AsyncSession):
"""Removing a non-existent association does not raise."""
post, (tag1,) = await self._setup(
db_session, "rm_author4", "rm4@test.com", "tag_rm_g"
)
tag2 = Tag(name="tag_rm_h")
db_session.add(tag2)
await db_session.flush()
# tag2 was never associated — should not raise
async with transaction(db_session):
await m2m_remove(db_session, post, Post.tags, tag2)
remaining = await self._load_tags(db_session, post)
assert len(remaining) == 1
@pytest.mark.anyio
async def test_non_m2m_raises_type_error(self, db_session: AsyncSession):
"""Passing a non-M2M relationship attribute raises TypeError."""
user = User(username="rm_author5", email="rm5@test.com")
db_session.add(user)
await db_session.flush()
role = Role(name="rm_type_err_role")
db_session.add(role)
await db_session.flush()
with pytest.raises(TypeError, match="Many-to-Many"):
await m2m_remove(db_session, user, User.role, role)
@pytest.mark.anyio
async def test_removes_composite_pk_related(self):
"""Composite-PK branch: DELETE uses tuple IN when related side has multi-col PK."""
engine = create_async_engine(DATABASE_URL, echo=False)
async with engine.begin() as conn:
await conn.run_sync(_LocalBase.metadata.create_all)
session_factory = async_sessionmaker(engine, expire_on_commit=False)
try:
async with session_factory() as session:
owner = _CompOwner()
item1 = _CompItem(group_id="g1", item_code="c1")
item2 = _CompItem(group_id="g1", item_code="c2")
session.add_all([owner, item1, item2])
await session.flush()
async with transaction(session):
await m2m_add(session, owner, _CompOwner.items, item1, item2)
async with transaction(session):
await m2m_remove(session, owner, _CompOwner.items, item1)
await session.commit()
async with session_factory() as verify:
result = await verify.execute(
select(_CompOwner)
.where(_CompOwner.id == owner.id)
.options(selectinload(_CompOwner.items))
)
loaded = result.scalar_one()
assert len(loaded.items) == 1
assert (loaded.items[0].group_id, loaded.items[0].item_code) == (
"g1",
"c2",
)
finally:
async with engine.begin() as conn:
await conn.run_sync(_LocalBase.metadata.drop_all)
await engine.dispose()
class TestM2MSet:
"""Tests for m2m_set helper."""
async def _load_tags(self, session: AsyncSession, post: Post) -> list[Tag]:
result = await session.execute(
select(Post).where(Post.id == post.id).options(selectinload(Post.tags))
)
return result.scalar_one().tags
@pytest.mark.anyio
async def test_replaces_existing_set(self, db_session: AsyncSession):
"""Replaces the full association set atomically."""
user = User(username="set_author1", email="set1@test.com")
db_session.add(user)
await db_session.flush()
post = Post(title="Post Set A", author_id=user.id)
tag1 = Tag(name="tag_set_a")
tag2 = Tag(name="tag_set_b")
tag3 = Tag(name="tag_set_c")
db_session.add_all([post, tag1, tag2, tag3])
await db_session.flush()
async with transaction(db_session):
await m2m_add(db_session, post, Post.tags, tag1, tag2)
async with transaction(db_session):
await m2m_set(db_session, post, Post.tags, tag3)
remaining = await self._load_tags(db_session, post)
assert len(remaining) == 1
assert remaining[0].id == tag3.id
@pytest.mark.anyio
async def test_clears_all_when_no_related(self, db_session: AsyncSession):
"""Passing no related instances clears all associations."""
user = User(username="set_author2", email="set2@test.com")
db_session.add(user)
await db_session.flush()
post = Post(title="Post Set B", author_id=user.id)
tag = Tag(name="tag_set_d")
db_session.add_all([post, tag])
await db_session.flush()
async with transaction(db_session):
await m2m_add(db_session, post, Post.tags, tag)
async with transaction(db_session):
await m2m_set(db_session, post, Post.tags)
remaining = await self._load_tags(db_session, post)
assert remaining == []
@pytest.mark.anyio
async def test_set_on_empty_then_populate(self, db_session: AsyncSession):
"""m2m_set works on a post with no existing associations."""
user = User(username="set_author3", email="set3@test.com")
db_session.add(user)
await db_session.flush()
post = Post(title="Post Set C", author_id=user.id)
tag1 = Tag(name="tag_set_e")
tag2 = Tag(name="tag_set_f")
db_session.add_all([post, tag1, tag2])
await db_session.flush()
async with transaction(db_session):
await m2m_set(db_session, post, Post.tags, tag1, tag2)
remaining = await self._load_tags(db_session, post)
assert {t.id for t in remaining} == {tag1.id, tag2.id}
@pytest.mark.anyio
async def test_non_m2m_raises_type_error(self, db_session: AsyncSession):
"""Passing a non-M2M relationship attribute raises TypeError."""
user = User(username="set_author4", email="set4@test.com")
db_session.add(user)
await db_session.flush()
role = Role(name="set_type_err_role")
db_session.add(role)
await db_session.flush()
with pytest.raises(TypeError, match="Many-to-Many"):
await m2m_set(db_session, user, User.role, role)
STATE_ATTR = "test_db_session"
class _FakeSession:
"""Records commit() calls into a shared event log."""
def __init__(self, events: list[str], *, in_txn: bool = True) -> None:
self.events = events
self._in_txn = in_txn
self.commits = 0
def in_transaction(self) -> bool:
return self._in_txn
async def commit(self) -> None:
self.events.append("COMMIT")
self.commits += 1
self._in_txn = False
async def _drive(app, scope, events: list[str]) -> list[str]:
"""Run an ASGI app, appending the response messages it emits to *events*
(shared with the fake session so commit/response ordering is captured)."""
async def receive(): # pragma: no cover - not exercised
return {"type": "http.disconnect"}
async def send(message) -> None:
events.append(message["type"])
await app(scope, receive, send)
return events
class TestCommitOrdering:
"""The commit must precede the forwarded response, and be skipped otherwise."""
@pytest.mark.anyio
async def test_commits_before_response_start(self):
events: list[str] = []
session = _FakeSession(events)
async def inner(scope, receive, send):
await send({"type": "http.response.start", "status": 200, "headers": []})
await send({"type": "http.response.body", "body": b"ok"})
app = _CommitOnResponseMiddleware(inner, state_attr=STATE_ATTR)
scope = {"type": "http", "state": {STATE_ATTR: session}}
result = await _drive(app, scope, events)
assert session.commits == 1
assert result == ["COMMIT", "http.response.start", "http.response.body"]
@pytest.mark.anyio
async def test_no_commit_when_not_in_transaction(self):
events: list[str] = []
session = _FakeSession(events, in_txn=False)
async def inner(scope, receive, send):
await send({"type": "http.response.start", "status": 200, "headers": []})
await send({"type": "http.response.body", "body": b""})
app = _CommitOnResponseMiddleware(inner, state_attr=STATE_ATTR)
scope = {"type": "http", "state": {STATE_ATTR: session}}
result = await _drive(app, scope, events)
assert session.commits == 0
assert result == ["http.response.start", "http.response.body"]
@pytest.mark.anyio
async def test_no_session_is_noop(self):
events: list[str] = []
async def inner(scope, receive, send):
await send({"type": "http.response.start", "status": 200, "headers": []})
await send({"type": "http.response.body", "body": b""})
app = _CommitOnResponseMiddleware(inner, state_attr=STATE_ATTR)
scope = {"type": "http", "state": {}}
result = await _drive(app, scope, events)
assert result == ["http.response.start", "http.response.body"]
@pytest.mark.anyio
async def test_non_http_scope_passes_through(self):
called = False
async def inner(scope, receive, send):
nonlocal called
called = True
async def receive(): # pragma: no cover - not exercised
return {"type": "lifespan.startup"}
async def send(message): # pragma: no cover - not exercised
return None
app = _CommitOnResponseMiddleware(inner, state_attr=STATE_ATTR)
await app({"type": "lifespan"}, receive, send)
assert called is True
class _ProbeMiddleware:
"""Outer middleware that records, at response start, whether a row created
in the request is already visible to a *separate* session."""
def __init__(self, app, *, session_maker, name: str, result: dict) -> None:
self.app = app
self.session_maker = session_maker
self.name = name
self.result = result
async def __call__(self, scope, receive, send):
async def send_wrapper(message):
if message["type"] == "http.response.start":
async with self.session_maker() as probe:
row = (
await probe.execute(select(Role).where(Role.name == self.name))
).scalar_one_or_none()
self.result["visible_at_start"] = row is not None
await send(message)
await self.app(scope, receive, send_wrapper)
def _build_app(db: Database) -> FastAPI:
"""A FastAPI app wired with the Database dependency and commit middleware."""
app = FastAPI()
@app.post("/roles")
async def create_role(
body: RoleCreate, session: AsyncSession = Depends(db)
) -> dict:
role = await RoleCrud.create(session, body)
return {"id": str(role.id), "name": role.name}
@app.post("/roles-then-boom")
async def create_then_raise(
body: RoleCreate, session: AsyncSession = Depends(db)
) -> dict:
await RoleCrud.create(session, body)
raise RuntimeError("boom after write")
@app.post("/two-roles")
async def create_two_roles(
body: RoleCreate, session: AsyncSession = Depends(db)
) -> dict:
# First write succeeds, second collides on the unique name and must
# take the whole request transaction down with it.
await RoleCrud.create(session, body)
await RoleCrud.create(session, body)
return {"ok": True}
@app.post("/roles-self-commit")
async def create_then_self_commit(
body: RoleCreate, session: AsyncSession = Depends(db)
) -> dict:
# Endpoint commits explicitly; the middleware must not double-commit or
# error — it finds no open transaction and no-ops.
role = await RoleCrud.create(session, body)
await session.commit()
return {"id": str(role.id), "name": role.name}
@app.get("/roles-stream/{name}")
async def stream_role(
name: str, session: AsyncSession = Depends(db)
) -> StreamingResponse:
# A write before the stream begins: the middleware commits it at
# response-start, before the generator runs.
await RoleCrud.create(session, RoleCreate(name=name))
async def gen():
# Read-only DB use during the stream, via the request session.
row = (
await session.execute(select(Role).where(Role.name == name))
).scalar_one()
yield f"data: {row.name}\n\n".encode()
return StreamingResponse(gen(), media_type="text/event-stream")
db.install(app)
return app
async def _row_exists(session_maker, name: str) -> bool:
async with session_maker() as session:
row = (
await session.execute(select(Role).where(Role.name == name))
).scalar_one_or_none()
return row is not None
class TestCommitIntegration:
@pytest.mark.anyio
async def test_write_is_committed(self, engine, session_maker):
app = _build_app(Database(engine=engine))
transport = ASGITransport(app=app)
async with AsyncClient(transport=transport, base_url="http://test") as client:
resp = await client.post("/roles", json={"name": "committed_role"})
assert resp.status_code == 200
assert await _row_exists(session_maker, "committed_role")
@pytest.mark.anyio
async def test_visible_at_response_start(self, engine, session_maker):
"""The write is visible to a separate session *before* the response is
sent — the read-after-write guarantee the middleware exists for."""
app = _build_app(Database(engine=engine))
result: dict = {}
app.add_middleware(
_ProbeMiddleware,
session_maker=session_maker,
name="probe_role",
result=result,
)
transport = ASGITransport(app=app)
async with AsyncClient(transport=transport, base_url="http://test") as client:
resp = await client.post("/roles", json={"name": "probe_role"})
assert resp.status_code == 200
assert result.get("visible_at_start") is True
@pytest.mark.anyio
async def test_error_rolls_back(self, engine, session_maker):
app = _build_app(Database(engine=engine))
transport = ASGITransport(app=app, raise_app_exceptions=False)
async with AsyncClient(transport=transport, base_url="http://test") as client:
resp = await client.post("/roles-then-boom", json={"name": "ghost_role"})
assert resp.status_code == 500
assert not await _row_exists(session_maker, "ghost_role")
@pytest.mark.anyio
async def test_explicit_commit_in_endpoint(self, engine, session_maker):
"""An endpoint that commits itself works: the middleware no-ops (no
double commit / error) and the write is persisted."""
app = _build_app(Database(engine=engine))
transport = ASGITransport(app=app)
async with AsyncClient(transport=transport, base_url="http://test") as client:
resp = await client.post("/roles-self-commit", json={"name": "self_commit"})
assert resp.status_code == 200
assert await _row_exists(session_maker, "self_commit")
@pytest.mark.anyio
async def test_streaming_response_coexists(self, engine, session_maker):
"""A read-only streaming endpoint works alongside the middleware: the
commit fires at stream start, the pre-stream write is committed, and the
generator can keep reading via the request session."""
app = _build_app(Database(engine=engine))
transport = ASGITransport(app=app)
async with AsyncClient(transport=transport, base_url="http://test") as client:
resp = await client.get("/roles-stream/streamed_role")
assert resp.status_code == 200
assert "data: streamed_role" in resp.text
# The write made before the stream began is durably committed.
assert await _row_exists(session_maker, "streamed_role")
@pytest.mark.anyio
async def test_multi_write_atomicity(self, engine, session_maker):
"""When the 2nd write fails, the 1st must roll back too (one txn)."""
app = _build_app(Database(engine=engine))
transport = ASGITransport(app=app, raise_app_exceptions=False)
async with AsyncClient(transport=transport, base_url="http://test") as client:
resp = await client.post("/two-roles", json={"name": "dup_role"})
assert resp.status_code >= 400
assert not await _row_exists(session_maker, "dup_role")