如何在Django的pytest(--no-migrations)创建表前安装PostgreSQL citext扩展？

我在Django项目中通过迁移可以正常安装PostgreSQL的citext扩展，但用pytest时加了--no-migrations选项，导致测试数据库创建时不执行迁移。现在pytest创建auth_user表时抛出错误：

django.db.utils.ProgrammingError: type "citext" does not exist

对应的SQL语句片段：

CREATE TABLE "auth_user" ("id" serial NOT NULL PRIMARY KEY, "password" varchar(128) NOT NULL, "last_login" timestamp ...T NULL, "is_active" boolean NOT NULL, "date_joined" timestamp with time zone NOT NULL, "email" citext NOT NULL UNIQUE)

我试过修改django_db_setup fixture，但没解决问题，代码如下：

@pytest.fixture(scope="session")
def django_db_setup(
    request,
    django_test_environment,
    django_db_blocker,
    django_db_use_migrations,
    django_db_keepdb,
    django_db_createdb,
    django_db_modify_db_settings,
):
    """Top level fixture to ensure test databases are available"""
    from django.test.utils import setup_databases, teardown_databases

    setup_databases_args = {}

    if not django_db_use_migrations:
        from pytest_django.fixtures import _disable_native_migrations
        _disable_native_migrations()

    if django_db_keepdb and not django_db_createdb:
        setup_databases_args["keepdb"] = True

    with django_db_blocker.unblock():
        from django.db import connection
        cursor = connection.cursor()
        cursor.execute("CREATE EXTENSION IF NOT EXISTS citext;")

        db_cfg = setup_databases(
            verbosity=request.config.option.verbose,
            interactive=False,
            **setup_databases_args
        )
    def teardown_database():
        with django_db_blocker.unblock():
            try:
                teardown_databases(db_cfg, verbosity=request.config.option.verbose)
            except Exception as exc:
                request.node.warn(
                    pytest.PytestWarning(
                        "Error when trying to teardown test databases: %r" % exc
                    )
                )

    if not django_db_keepdb:
        request.addfinalizer(teardown_database)

问题原因 #

你执行CREATE EXTENSION的时机不对——调用setup_databases()之前执行时，测试数据库还未创建，此时连接的是默认数据库而非测试库，导致扩展没安装到正确的位置。

解决方案 #

方法一：修正django_db_setup fixture的执行顺序 #

调整逻辑，先创建测试数据库，再连接测试库安装扩展：

@pytest.fixture(scope="session")
def django_db_setup(
    request,
    django_test_environment,
    django_db_blocker,
    django_db_use_migrations,
    django_db_keepdb,
    django_db_createdb,
    django_db_modify_db_settings,
):
    from django.test.utils import setup_databases, teardown_databases

    setup_databases_args = {}

    if not django_db_use_migrations:
        from pytest_django.fixtures import _disable_native_migrations
        _disable_native_migrations()

    if django_db_keepdb and not django_db_createdb:
        setup_databases_args["keepdb"] = True

    with django_db_blocker.unblock():
        # 先创建测试数据库
        db_cfg = setup_databases(
            verbosity=request.config.option.verbose,
            interactive=False,
            **setup_databases_args
        )
        # 连接到测试数据库并安装citext扩展
        from django.db import connections
        for conn_name in connections:
            conn = connections[conn_name]
            with conn.cursor() as cursor:
                cursor.execute("CREATE EXTENSION IF NOT EXISTS citext;")

    def teardown_database():
        with django_db_blocker.unblock():
            try:
                teardown_databases(db_cfg, verbosity=request.config.option.verbose)
            except Exception as exc:
                request.node.warn(
                    pytest.PytestWarning(
                        "Error when trying to teardown test databases: %r" % exc
                    )
                )

    if not django_db_keepdb:
        request.addfinalizer(teardown_database)

方法二：自定义测试数据库创建类（更持久） #

通过修改Django的数据库创建逻辑，让测试数据库自动安装扩展，无需修改pytest配置：

1. 在项目测试目录下创建工具文件（比如tests/utils.py）：

from django.contrib.gis.db.backends.postgis.creation import PostGISDatabaseCreation

class CustomPostGISDatabaseCreation(PostGISDatabaseCreation):
    def create_test_db(self, *args, **kwargs):
        # 先调用父类创建数据库
        db_name = super().create_test_db(*args, **kwargs)
        # 安装citext扩展
        with self.connection.cursor() as cursor:
            cursor.execute("CREATE EXTENSION IF NOT EXISTS citext;")
        return db_name

2. 在项目settings.py中指定自定义创建类：

DATABASES = {
    'default': {
        'ENGINE': 'django.contrib.gis.db.backends.postgis',
        # ...其他数据库配置...
        'TEST': {
            'CREATION_CLASS': 'tests.utils.CustomPostGISDatabaseCreation',
        },
    }
}

这样无论是否使用--no-migrations，测试数据库创建时都会自动安装citext扩展。

内容的提问来源于stack exchange，提问作者Alexandr Zayets