如何编写单元测试验证PostgreSQL数据库创建函数？

一、测试环境准备 #

• 单独搭建测试用PostgreSQL实例（比如本地临时服务或Docker容器），绝对不要直接操作生产环境数据库。
• 配置独立的测试连接参数（用户、密码、端口），与生产环境完全隔离。

二、核心测试用例设计 #

1. 正常创建数据库场景 #

• 传入合法的数据库名称（如test_valid_db_001），调用函数后验证数据库是否存在：
  • 连接到默认的postgres库，执行SELECT 1 FROM pg_database WHERE datname = 'test_valid_db_001'，检查是否有返回结果。

2. 重复创建已存在数据库场景 #

• 先手动创建目标数据库，再调用函数传入相同名称，验证是否抛出psycopg2.OperationalError（PostgreSQL数据库已存在的标准异常）。

3. 非法数据库名称场景 #

• 传入包含特殊字符、空格或PostgreSQL保留字的名称（如my db、user、test-db!），验证函数是否抛出对应异常（语法错误或操作错误）。
• 注：你的函数未对数据库名称加引号，大写名称会被自动转小写，这类场景也可纳入测试。

4. 无效参数场景 #

• 传入空字符串、None或全空格字符串，验证函数是否能正确处理（比如抛出参数错误或数据库执行报错）。

三、测试实现的两种方案 #

方案1：真实数据库集成测试 #

优点：完全模拟实际执行流程，覆盖数据库交互的真实逻辑。
步骤：

1. 测试前清理：删除可能存在的测试数据库，避免干扰。
2. 调用目标函数执行创建操作。
3. 执行验证逻辑，确认结果符合预期。
4. 测试后清理：删除测试生成的数据库，释放资源。

示例代码（基于pytest）：

import psycopg2
import pytest
from your_module import create_database

def test_create_database_success():
    test_db = "test_success_db"
    # 前置清理
    conn = psycopg2.connect(database="postgres", user='postgres', password='password', host='127.0.0.1', port='5432')
    conn.autocommit = True
    conn.cursor().execute(f"DROP DATABASE IF EXISTS {test_db}")
    conn.close()

    # 执行创建
    create_database(test_db)

    # 验证存在
    conn = psycopg2.connect(database="postgres", user='postgres', password='password', host='127.0.0.1', port='5432')
    cursor = conn.cursor()
    cursor.execute(f"SELECT 1 FROM pg_database WHERE datname = '{test_db}'")
    assert cursor.fetchone() is not None
    conn.close()

def test_create_existing_database():
    test_db = "test_existing_db"
    # 前置创建数据库
    conn = psycopg2.connect(database="postgres", user='postgres', password='password', host='127.0.0.1', port='5432')
    conn.autocommit = True
    conn.cursor().execute(f"CREATE DATABASE IF NOT EXISTS {test_db}")
    conn.close()

    # 验证异常
    with pytest.raises(psycopg2.OperationalError):
        create_database(test_db)

    # 后置清理
    conn = psycopg2.connect(database="postgres", user='postgres', password='password', host='127.0.0.1', port='5432')
    conn.autocommit = True
    conn.cursor().execute(f"DROP DATABASE IF EXISTS {test_db}")
    conn.close()

方案2：Mock隔离单元测试 #

优点：无需真实数据库，测试速度快，专注于函数逻辑本身。
步骤：

1. Mockpsycopg2.connect方法，返回模拟的连接对象。
2. Mock连接对象的cursor方法，返回模拟的游标对象。
3. 验证游标是否执行了正确的SQL语句，连接参数、状态是否符合预期。

示例代码（基于unittest.mock）：

from unittest.mock import patch, MagicMock
from your_module import create_database

def test_create_database_with_mock():
    test_db = "test_mock_db"
    with patch('psycopg2.connect') as mock_connect:
        # 模拟连接与游标
        mock_conn = MagicMock()
        mock_connect.return_value = mock_conn
        mock_cursor = MagicMock()
        mock_conn.cursor.return_value = mock_cursor

        # 执行函数
        create_database(test_db)

        # 验证调用逻辑
        mock_connect.assert_called_once_with(
            database="postgres", user='postgres', password='password', host='127.0.0.1', port='5432'
        )
        assert mock_conn.autocommit is True
        mock_cursor.execute.assert_called_once_with(f'''CREATE database {test_db}''')
        mock_conn.close.assert_called_once()

四、额外注意事项 #

• 你的函数使用f-string拼接SQL，存在SQL注入风险（比如传入test_db; DROP DATABASE other_db;），建议测试这类恶意输入场景，后续可对数据库名称做校验（仅允许字母、数字、下划线）。
• 确保测试用例独立，每个用例执行前后清理测试资源，避免用例间互相干扰。
• 可使用pytest fixture封装数据库连接、清理逻辑，减少代码重复。

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