1import re
  2
  3import pytest
  4
  5from plain.models.otel import suppress_db_tracing
  6from plain.signals import request_finished, request_started
  7
  8from .. import transaction
  9from ..backends.base.base import BaseDatabaseWrapper
 10from ..db import close_old_connections, db_connection
 11from .utils import (
 12    setup_database,
 13    teardown_database,
 14)
 15
 16
 17@pytest.fixture(autouse=True)
 18def _db_disabled():
 19    """
 20    Every test should use this fixture by default to prevent
 21    access to the normal database.
 22    """
 23
 24    def cursor_disabled(self):
 25        pytest.fail("Database access not allowed without the `db` fixture")
 26
 27    BaseDatabaseWrapper._enabled_cursor = BaseDatabaseWrapper.cursor
 28    BaseDatabaseWrapper.cursor = cursor_disabled
 29
 30    yield
 31
 32    BaseDatabaseWrapper.cursor = BaseDatabaseWrapper._enabled_cursor
 33
 34
 35@pytest.fixture(scope="session")
 36def setup_db(request):
 37    """
 38    This fixture is called automatically by `db`,
 39    so a test database will only be setup if the `db` fixture is used.
 40    """
 41    verbosity = request.config.option.verbose
 42
 43    # Set up the test db across the entire session
 44    _old_db_name = setup_database(verbosity=verbosity)
 45
 46    # Keep connections open during request client / testing
 47    request_started.disconnect(close_old_connections)
 48    request_finished.disconnect(close_old_connections)
 49
 50    yield
 51
 52    # Put the signals back...
 53    request_started.connect(close_old_connections)
 54    request_finished.connect(close_old_connections)
 55
 56    # When the test session is done, tear down the test db
 57    teardown_database(_old_db_name, verbosity=verbosity)
 58
 59
 60@pytest.fixture
 61def db(setup_db, request):
 62    if "isolated_db" in request.fixturenames:
 63        pytest.fail("The 'db' and 'isolated_db' fixtures cannot be used together")
 64
 65    # Set .cursor() back to the original implementation to unblock it
 66    BaseDatabaseWrapper.cursor = BaseDatabaseWrapper._enabled_cursor
 67
 68    if not db_connection.features.supports_transactions:
 69        pytest.fail("Database does not support transactions")
 70
 71    with suppress_db_tracing():
 72        atomic = transaction.atomic()
 73        atomic._from_testcase = True  # TODO remove this somehow?
 74        atomic.__enter__()
 75
 76    yield
 77
 78    with suppress_db_tracing():
 79        if (
 80            db_connection.features.can_defer_constraint_checks
 81            and not db_connection.needs_rollback
 82            and db_connection.is_usable()
 83        ):
 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):
 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    BaseDatabaseWrapper.cursor = BaseDatabaseWrapper._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)