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 ..backends.base.base import BaseDatabaseWrapper
 14from ..db import close_old_connections, db_connection
 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    BaseDatabaseWrapper._enabled_cursor = BaseDatabaseWrapper.cursor  # type: ignore[attr-defined]
 32    BaseDatabaseWrapper.cursor = cursor_disabled  # type: ignore[method-assign]
 33
 34    yield
 35
 36    BaseDatabaseWrapper.cursor = BaseDatabaseWrapper._enabled_cursor  # type: ignore[method-assign]
 37
 38
 39@pytest.fixture(scope="session")
 40def setup_db(request: Any) -> Generator[None, None, None]:
 41    """
 42    This fixture is called automatically by `db`,
 43    so a test database will only be setup if the `db` fixture is used.
 44    """
 45    verbosity = request.config.option.verbose
 46
 47    # Set up the test db across the entire session
 48    _old_db_name = setup_database(verbosity=verbosity)
 49
 50    # Keep connections open during request client / testing
 51    request_started.disconnect(close_old_connections)
 52    request_finished.disconnect(close_old_connections)
 53
 54    yield
 55
 56    # Put the signals back...
 57    request_started.connect(close_old_connections)
 58    request_finished.connect(close_old_connections)
 59
 60    # When the test session is done, tear down the test db
 61    teardown_database(_old_db_name, verbosity=verbosity)
 62
 63
 64@pytest.fixture
 65def db(setup_db: Any, request: Any) -> Generator[None, None, None]:
 66    if "isolated_db" in request.fixturenames:
 67        pytest.fail("The 'db' and 'isolated_db' fixtures cannot be used together")
 68
 69    # Set .cursor() back to the original implementation to unblock it
 70    BaseDatabaseWrapper.cursor = BaseDatabaseWrapper._enabled_cursor  # type: ignore[method-assign]
 71
 72    if not db_connection.features.supports_transactions:
 73        pytest.fail("Database does not support transactions")
 74
 75    with suppress_db_tracing():
 76        atomic = transaction.atomic()
 77        atomic._from_testcase = True  # type: ignore[attr-defined]  # TODO remove this somehow?
 78        atomic.__enter__()
 79
 80    yield
 81
 82    with suppress_db_tracing():
 83        if (
 84            db_connection.features.can_defer_constraint_checks
 85            and not db_connection.needs_rollback
 86            and db_connection.is_usable()
 87        ):
 88            db_connection.check_constraints()
 89
 90        db_connection.set_rollback(True)
 91        atomic.__exit__(None, None, None)
 92
 93        db_connection.close()
 94
 95
 96@pytest.fixture
 97def isolated_db(request: Any) -> Generator[None, None, None]:
 98    """
 99    Create and destroy a unique test database for each test, using a prefix
100    derived from the test function name to ensure isolation from the default
101    test database.
102    """
103    if "db" in request.fixturenames:
104        pytest.fail("The 'db' and 'isolated_db' fixtures cannot be used together")
105    # Set .cursor() back to the original implementation to unblock it
106    BaseDatabaseWrapper.cursor = BaseDatabaseWrapper._enabled_cursor  # type: ignore[method-assign]
107
108    verbosity = 1
109
110    # Derive a safe prefix from the test function name
111    raw_name = request.node.name
112    prefix = re.sub(r"[^0-9A-Za-z_]+", "_", raw_name)
113
114    # Set up a fresh test database for this test, using the prefix
115    _old_db_name = setup_database(verbosity=verbosity, prefix=prefix)
116
117    yield
118
119    # Tear down the test database created for this test
120    teardown_database(_old_db_name, verbosity=verbosity)