Search code examples
djangopostgresqlpytest-django

Setting up pytest django database


I need to run this SQL after the test database is created but before creating any tables:

`CREATE EXTENSION pg_trgm;
CREATE EXTENSION btree_gin;

I know that I can somehow use django_db_setup fixture to do this but I can't figure out how.
I use

addopts = --no-migrations

in my pytest.ini file so this can't be done using migrations

It should be something like this:

@pytest.fixture(scope='session')
def django_db_setup():
    # Some code to create database
        with connection.cursor() as cursor:
            cursor.execute('CREATE EXTENSION IF NOT EXISTS pg_trgm;')
            cursor.execute('CREATE EXTENSION IF NOT EXISTS btree_gin;')
    # Code for creating tables

Solution

  • @pytest.fixture(scope="session")
    def django_migration_disabler() -> None:
        """Disable migrations when running django tests.
    
        This copies/alters the behavior of pytest_django.fixtures._disable_migrations,
        which is called when pytest is invoked with --no-migrations.
    
        See: https://github.com/pytest-dev/pytest-django/blob/v4.5.2/pytest_django/fixtures.py#L260
    
        We do this instead of invoking with --no-migrations because constructing the
        database without migrations fails to create necessary postgres extensions like
        citext and uuid-ossp.
    
        We then override the django_db_setup fixture and ensure that this
        fixture is called before the parent django_db_setup fixture.
        """
        from django.conf import settings
        from django.core.management.commands import migrate
        from django.db import connections
    
        class DisableMigrations:
            def __contains__(self, item: str) -> bool:
                return True
    
            def __getitem__(self, item: str) -> None:
                return None
    
        settings.MIGRATION_MODULES = DisableMigrations()
    
        class MigrateSilentCommand(migrate.Command):
            def handle(self, *args, **options):
                options["verbosity"] = 0
                database = options["database"]
                connection = connections[database]
                with connection.cursor() as cursor:
                    cursor.execute('CREATE EXTENSION IF NOT EXISTS pg_trgm;')
                    cursor.execute('CREATE EXTENSION IF NOT EXISTS btree_gin;')
                return super().handle(*args, **options)
    
        migrate.Command = MigrateSilentCommand  # type: ignore
    
    
    @pytest.fixture(scope="session")
    def django_db_use_migrations() -> bool:
        """Force pytest-django to use migrations.
    
        This is necessary because we disable migrations in the django_migration_disabler
        and the existing pytest-django mechanisms would override our disabling mechanism
        with their own, which would fail to create the necessary extensions in postgres.
        """
        return True
    
    
    @pytest.fixture(scope="session")
    def django_db_setup(django_migration_disabler, django_db_setup, django_db_blocker):
        """Override django_db_setup ensuring django_migration_disabler is loaded."""
        pass