如何在Django pytest测试中适配自定义pg_trgm扩展数据库脚本？

可以编写自定义pytest fixture，在测试数据库初始化完成后执行创建pg_trgm扩展的SQL语句，以下是具体实现：

1. 编写自定义Fixture #

在测试目录的conftest.py文件中添加以下代码：

import pytest
from django.db import connection

@pytest.fixture(scope="session")
def setup_pg_trgm_extension(django_db_setup, django_db_blocker):
    # 解锁数据库操作上下文，安全执行自定义SQL
    with django_db_blocker.unblock():
        with connection.cursor() as cursor:
            # 用IF NOT EXISTS避免重复创建报错
            cursor.execute("CREATE EXTENSION IF NOT EXISTS pg_trgm;")

2. 在测试中使用Fixture #

方式一：按需引入 #

在需要用到pg_trgm扩展的测试用例中，直接声明依赖该fixture：

def test_trigram_based_query(setup_pg_trgm_extension):
    # 编写测试逻辑，比如创建MyModel实例并验证trigram索引的查询效果
    pass

方式二：全局自动应用 #

如果希望所有测试都自动加载pg_trgm扩展，可修改fixture为自动执行：

@pytest.fixture(scope="session", autouse=True)
def setup_pg_trgm_extension(django_db_setup, django_db_blocker):
    with django_db_blocker.unblock():
        with connection.cursor() as cursor:
            cursor.execute("CREATE EXTENSION IF NOT EXISTS pg_trgm;")

原理说明 #

• django_db_setup：依赖这个fixture确保测试数据库已经完成初始化
• django_db_blocker：处理pytest-django的数据库锁定机制，保证自定义SQL能安全执行
• CREATE EXTENSION IF NOT EXISTS：避免重复创建扩展引发的错误，兼容多次执行场景

额外注意 #

• 会话级别（scope="session"）的fixture只会执行一次，足够覆盖整个测试会话的需求；若需每次测试都重新创建扩展，可将scope改为function
• 如果测试使用了事务型数据库fixture（默认的django_db），会话级别的初始化不会被事务回滚影响，扩展会一直存在于测试数据库中

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