如何在pytest中创建PostgreSQL数据库会话解决连接过多问题？

核心思路：复用数据库连接 #

不要在每个查询函数里单独创建和关闭连接，而是通过pytest的fixture管理一个全局（或指定作用域内）的连接，所有查询函数共享该连接，测试结束后统一关闭。

具体实现步骤 #

1. 定义数据库连接fixture #

使用pytest的session或module作用域fixture，确保整个测试会话（或模块）仅创建一次连接：

import pytest
import pg8000
import constants

@pytest.fixture(scope="session")
def db_connection():
    conn = pg8000.connect(**constants.AUTOTESTS_DB_CREDS)
    yield conn
    # 测试会话结束后统一关闭连接
    conn.close()

2. 改造查询函数，复用连接 #

修改查询函数，接收外部传入的连接，不再单独创建连接：

def get_expected_status_code_API01(db_connection, role, incident_type="a"):
    cursor = db_connection.cursor()
    try:
        cursor.execute(
            "SELECT status_code FROM pytest.api01_roles_and_status_codes WHERE role=%s AND incident_type=%s",
            [role, incident_type]
        )
        status_code = cursor.fetchall()[0][0]
        return status_code
    finally:
        # 用完游标立即关闭，连接保留复用
        cursor.close()

3. 在测试用例或fixture中使用 #

在需要调用查询函数的地方，传入db_connection fixture：

def test_api01_status_code(db_connection):
    status_code = get_expected_status_code_API01(db_connection, "admin")
    assert status_code == 200

额外优化建议 #

• 游标管理：每次查询后及时关闭游标，避免不必要的资源占用
• 事务隔离：如果测试涉及数据修改，可添加函数级fixture实现事务回滚，保证测试环境干净：

  @pytest.fixture(scope="function")
  def db_session(db_connection):
    db_connection.autocommit = False
    yield db_connection
    # 每个测试用例结束后回滚事务
    db_connection.rollback()
  

  之后用db_session替代db_connection传入查询函数即可。
• 连接池方案：若测试会话长或需并发执行，可使用连接池（如pg8000内置连接池或SQLAlchemy连接池）进一步优化连接分配。

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