1import pytest
  2
  3from plain.signals import request_finished, request_started
  4
  5from .. import transaction
  6from ..backends.base.base import BaseDatabaseWrapper
  7from ..db import close_old_connections, connections
  8from .utils import (
  9    setup_databases,
 10    teardown_databases,
 11)
 12
 13
 14@pytest.fixture(autouse=True)
 15def _db_disabled():
 16    """
 17    Every test should use this fixture by default to prevent
 18    access to the normal database.
 19    """
 20
 21    def cursor_disabled(self):
 22        pytest.fail("Database access not allowed without the `db` fixture")
 23
 24    BaseDatabaseWrapper._old_cursor = BaseDatabaseWrapper.cursor
 25    BaseDatabaseWrapper.cursor = cursor_disabled
 26
 27    yield
 28
 29    BaseDatabaseWrapper.cursor = BaseDatabaseWrapper._old_cursor
 30
 31
 32@pytest.fixture(scope="session")
 33def setup_db(request):
 34    """
 35    This fixture is called automatically by `db`,
 36    so a test database will only be setup if the `db` fixture is used.
 37    """
 38    verbosity = request.config.option.verbose
 39
 40    # Set up the test db across the entire session
 41    _old_db_config = setup_databases(verbosity=verbosity)
 42
 43    # Keep connections open during request client / testing
 44    request_started.disconnect(close_old_connections)
 45    request_finished.disconnect(close_old_connections)
 46
 47    yield _old_db_config
 48
 49    # Put the signals back...
 50    request_started.connect(close_old_connections)
 51    request_finished.connect(close_old_connections)
 52
 53    # When the test session is done, tear down the test db
 54    teardown_databases(_old_db_config, verbosity=verbosity)
 55
 56
 57@pytest.fixture()
 58def db(setup_db):
 59    # Set .cursor() back to the original implementation
 60    BaseDatabaseWrapper.cursor = BaseDatabaseWrapper._old_cursor
 61
 62    # Keep track of the atomic blocks so we can roll them back
 63    atomics = {}
 64
 65    for connection in connections.all():
 66        # By default we use transactions to rollback changes,
 67        # so we need to ensure the database supports transactions
 68        if not connection.features.supports_transactions:
 69            pytest.fail("Database does not support transactions")
 70
 71        # Clear the queries log before each test?
 72        # connection.queries_log.clear()
 73
 74        atomic = transaction.atomic(using=connection.alias)
 75        atomic._from_testcase = True  # TODO remove this somehow?
 76        atomic.__enter__()
 77        atomics[connection] = atomic
 78
 79    yield setup_db
 80
 81    for connection, atomic in atomics.items():
 82        if (
 83            connection.features.can_defer_constraint_checks
 84            and not connection.needs_rollback
 85            and connection.is_usable()
 86        ):
 87            connection.check_constraints()
 88
 89        transaction.set_rollback(True, using=connection.alias)
 90        atomic.__exit__(None, None, None)
 91
 92        connection.close()
 93
 94
 95# @pytest.fixture(scope="function")
 96# def transactional_db(request, _plain_db_setup):
 97#     BaseDatabaseWrapper.cursor = BaseDatabaseWrapper._cursor
 98
 99#     yield
100
101#     # Flush databases instead of rolling back transactions...
102
103#     connections.close_all()