调用subprocess执行listen命令后Django测试数据库未自动删除

背景 #

我们用Django-pgpubsub替代Django signals或Celery，通过运行python manage.py listen监听PostgreSQL触发器，当Shipment模型的status字段变为特定值时触发邮件通知。集成测试中通过subprocess启动独立进程运行该监听命令，测试能正常通过，但测试数据库无法被pytest-django自动清理——即使日志显示Destroying test database for alias 'default' ('test_my_project_database')，数据库实际仍保留，后续测试会因数据库重复存在报错。目前用--reuse-db参数能规避问题，但不想依赖这个临时方案。

问题根源 #

测试中启动的listen进程会持续持有PostgreSQL数据库连接，pytest-django尝试销毁数据库时，该连接仍处于活跃状态。PostgreSQL不允许删除有活跃连接的数据库，哪怕调用了listen_process.terminate()，进程也可能没来得及正常关闭连接就被终止，导致数据库无法被删除。

解决方案 #

1. 确保监听进程彻底退出并释放连接 #

不要只用terminate()，改用“先终止+超时等待+强制杀死”的组合逻辑，确保进程完全退出：

# 替换原代码中的listen_process.terminate()
listen_process.terminate()
try:
    # 等待进程正常退出，设置超时时间
    listen_process.wait(timeout=5)
except subprocess.TimeoutExpired:
    # 超时则强制杀死进程
    listen_process.kill()
    listen_process.wait()

2. 动态获取测试数据库名称，避免硬编码 #

原代码中硬编码测试数据库名test_{config('POSTGRES_DB')}，但pytest-django未启用--reuse-db时，可能生成带随机后缀的测试数据库。正确做法是从Django当前数据库配置中获取名称：

from django.db import connections

# 获取当前测试数据库的真实名称
db_name = connections['default'].settings_dict['NAME']
env.update(
    {
        "DATABASE_URL":
            f"postgresql://{config('POSTGRES_USER')}:{config('POSTGRES_PASSWORD')}@127.0.0.1:5432/{db_name}",
    }
)

3. 用pytest fixture统一管理监听进程 #

把启动/关闭监听进程的逻辑封装成fixture，确保每个测试后都能可靠清理：

import pytest
import subprocess
import os
import time
from django.conf import settings
from decouple import config
from django.db import connections

@pytest.fixture(scope="function")
def pgpubsub_listen_process():
    db_name = connections['default'].settings_dict['NAME']
    listen_command = [
        "python",
        str(settings.ROOT_DIR / "manage.py"),
        "listen",
    ]
    env = os.environ.copy()
    env.update(
        {
            "DATABASE_URL":
                f"postgresql://{config('POSTGRES_USER')}:{config('POSTGRES_PASSWORD')}@127.0.0.1:5432/{db_name}",
        }
    )
    listen_process = subprocess.Popen(listen_command, env=env)
    time.sleep(2)  # 等待进程启动完成
    yield listen_process
    
    # 测试结束后清理进程
    listen_process.terminate()
    try:
        listen_process.wait(timeout=5)
    except subprocess.TimeoutExpired:
        listen_process.kill()
        listen_process.wait()

测试函数可以简化为：

@pytest.mark.django_db(transaction=True)
def test_shipment_status_notifications_with_listen_command_subprocess(testuser, pgpubsub_listen_process):
    user = testuser
    # 测试数据准备
    notification_settings = NotificationSettings.objects.get(user=user)
    notification_settings.failed_attempt = True
    notification_settings.save()
    cust_email = "test@test.com"
    customer = Customer.objects.create(email=cust_email)
    order = Order.objects.create(customer=customer)
    shipment = Shipment.objects.create(user=user, order=order)

    # 触发状态变更
    shipment.status = Shipment.FAILED_ATTEMPT
    shipment.save()

    time.sleep(2)  # 等待触发器处理完成

    existing_notifications = Notification.objects.filter(
        type=Notification.Type.FAILED_DELIVERY_ATTEMPT,
        shipment=shipment,
        email=cust_email,
    )
    assert existing_notifications.count() == 1

4. 极端情况：强制关闭数据库所有连接 #

如果以上方法仍无效，可以在测试清理阶段强制关闭目标数据库的所有活跃连接：

import psycopg2
from psycopg2 import sql

def close_db_connections(db_name):
    # 连接到postgres系统数据库（而非测试数据库）
    conn = psycopg2.connect(
        dbname="postgres",
        user=config('POSTGRES_USER'),
        password=config('POSTGRES_PASSWORD'),
        host="127.0.0.1",
        port="5432"
    )
    conn.autocommit = True
    cur = conn.cursor()
    # 终止所有连接到目标测试库的进程
    cur.execute(sql.SQL("""
        SELECT pg_terminate_backend(pid)
        FROM pg_stat_activity
        WHERE datname = %s AND pid <> pg_backend_pid();
    """), [db_name])
    cur.close()
    conn.close()

然后在fixture的清理阶段调用该函数：

@pytest.fixture(scope="function")
def pgpubsub_listen_process():
    # 启动进程逻辑...
    yield listen_process
    # 先终止进程
    listen_process.terminate()
    try:
        listen_process.wait(timeout=5)
    except subprocess.TimeoutExpired:
        listen_process.kill()
        listen_process.wait()
    # 强制关闭所有数据库连接
    db_name = connections['default'].settings_dict['NAME']
    close_db_connections(db_name)

验证方式 #

运行测试后，查看PostgreSQL数据库列表，确认测试数据库已被自动删除，后续测试不会因数据库重复存在报错。

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