1from __future__ import annotations
  2
  3import re
  4from collections.abc import Generator
  5from typing import Any
  6
  7import pytest
  8
  9from plain.models.otel import suppress_db_tracing
 10from plain.signals import request_finished, request_started
 11
 12from .. import transaction
 13from ..db import close_old_connections, db_connection
 14from ..postgres.wrapper import DatabaseWrapper
 15from .utils import (
 16    setup_database,
 17    teardown_database,
 18)
 19
 20
 21@pytest.fixture(autouse=True)
 22def _db_disabled() -> Generator[None, None, None]:
 23    """
 24    Every test should use this fixture by default to prevent
 25    access to the normal database.
 26    """
 27
 28    def cursor_disabled(self: Any) -> None:
 29        pytest.fail("Database access not allowed without the `db` fixture")
 30
 31    # Save original cursor method and replace with disabled version
 32    setattr(DatabaseWrapper, "_enabled_cursor", DatabaseWrapper.cursor)
 33    DatabaseWrapper.cursor = cursor_disabled  # type: ignore[assignment]
 34
 35    yield
 36
 37    # Restore original cursor method
 38    DatabaseWrapper.cursor = getattr(DatabaseWrapper, "_enabled_cursor")
 39
 40
 41@pytest.fixture(scope="session")
 42def setup_db(request: Any) -> Generator[None, None, None]:
 43    """
 44    This fixture is called automatically by `db`,
 45    so a test database will only be setup if the `db` fixture is used.
 46    """
 47    verbosity = request.config.option.verbose
 48
 49    # Set up the test db across the entire session
 50    _old_db_name = setup_database(verbosity=verbosity)
 51
 52    # Keep connections open during request client / testing
 53    request_started.disconnect(close_old_connections)
 54    request_finished.disconnect(close_old_connections)
 55
 56    yield
 57
 58    # Put the signals back...
 59    request_started.connect(close_old_connections)
 60    request_finished.connect(close_old_connections)
 61
 62    # When the test session is done, tear down the test db
 63    teardown_database(_old_db_name, verbosity=verbosity)
 64
 65
 66@pytest.fixture
 67def db(setup_db: Any, request: Any) -> Generator[None, None, None]:
 68    if "isolated_db" in request.fixturenames:
 69        pytest.fail("The 'db' and 'isolated_db' fixtures cannot be used together")
 70
 71    # Set .cursor() back to the original implementation to unblock it
 72    DatabaseWrapper.cursor = getattr(DatabaseWrapper, "_enabled_cursor")
 73
 74    with suppress_db_tracing():
 75        atomic = transaction.atomic()
 76        atomic._from_testcase = True
 77        atomic.__enter__()
 78
 79    yield
 80
 81    with suppress_db_tracing():
 82        # PostgreSQL can defer constraint checks
 83        if not db_connection.needs_rollback and db_connection.is_usable():
 84            db_connection.check_constraints()
 85
 86        db_connection.set_rollback(True)
 87        atomic.__exit__(None, None, None)
 88
 89        db_connection.close()
 90
 91
 92@pytest.fixture
 93def isolated_db(request: Any) -> Generator[None, None, None]:
 94    """
 95    Create and destroy a unique test database for each test, using a prefix
 96    derived from the test function name to ensure isolation from the default
 97    test database.
 98    """
 99    if "db" in request.fixturenames:
100        pytest.fail("The 'db' and 'isolated_db' fixtures cannot be used together")
101    # Set .cursor() back to the original implementation to unblock it
102    DatabaseWrapper.cursor = getattr(DatabaseWrapper, "_enabled_cursor")
103
104    verbosity = 1
105
106    # Derive a safe prefix from the test function name
107    raw_name = request.node.name
108    prefix = re.sub(r"[^0-9A-Za-z_]+", "_", raw_name)
109
110    # Set up a fresh test database for this test, using the prefix
111    _old_db_name = setup_database(verbosity=verbosity, prefix=prefix)
112
113    yield
114
115    # Tear down the test database created for this test
116    teardown_database(_old_db_name, verbosity=verbosity)