如何对数据库连接函数编写测试用例？Flask场景问题解析

一、测试思路验证 #

你的思路完全正确：单元测试通过Mock隔离外部数据库依赖，仅验证函数逻辑；集成测试则搭建真实测试数据库，验证完整连接与交互流程。

二、单元测试示例（Mock pymssql.connect） #

使用unittest.mock模拟数据库连接，重点验证环境变量读取、连接参数传递、自动提交设置及连接关闭逻辑：

import unittest
from unittest.mock import patch, MagicMock
import os
from your_module import db_connection  # 替换为实际模块路径

class TestDbConnection(unittest.TestCase):
    def setUp(self):
        # 预置测试环境变量
        os.environ['TEST_SERVER'] = 'test-server'
        os.environ['TEST_DATABASE'] = 'test-db'
        os.environ['TEST_USERNAME'] = 'test-user'
        os.environ['TEST_PASSWORD'] = 'test-pass'

    def tearDown(self):
        # 清理环境变量
        for key in ['TEST_SERVER', 'TEST_DATABASE', 'TEST_USERNAME', 'TEST_PASSWORD']:
            if key in os.environ:
                del os.environ[key]

    @patch('pymssql.connect')
    def test_db_connection_flow(self, mock_connect):
        # 模拟数据库连接对象
        mock_conn = MagicMock()
        mock_connect.return_value = mock_conn

        # 执行上下文管理器
        with db_connection('TEST') as conn:
            self.assertEqual(conn, mock_conn)
            # 验证自动提交配置
            mock_conn.autocommit.assert_called_once_with(True)
            # 验证连接参数是否正确传递（注意：原函数中timeout变量未定义，需确保实际代码已声明）
            mock_connect.assert_called_once_with(
                server='test-server',
                database='test-db',
                user='test-user',
                password='test-pass',
                timeout=timeout
            )

        # 验证连接最终被关闭
        mock_conn.close.assert_called_once()

三、集成测试示例（用Fixture搭建测试数据库） #

借助pytest的fixture创建临时测试数据库（建议用Docker快速搭建SQL Server测试环境）：

import pytest
import os
import pymssql
from your_module import db_connection

# 测试数据库基础配置
TEST_DB_CONFIG = {
    'server': 'localhost',
    'database': 'test_temp_db',
    'username': 'sa',
    'password': 'your-test-password'
}

@pytest.fixture(scope='module')
def test_db():
    # 1. 创建测试数据库（若不存在）
    conn = pymssql.connect(
        server=TEST_DB_CONFIG['server'],
        database='master',
        user=TEST_DB_CONFIG['username'],
        password=TEST_DB_CONFIG['password']
    )
    cursor = conn.cursor()
    cursor.execute(f"IF NOT EXISTS (SELECT * FROM sys.databases WHERE name = '{TEST_DB_CONFIG['database']}') CREATE DATABASE {TEST_DB_CONFIG['database']}")
    conn.commit()
    conn.close()

    # 2. 设置环境变量
    os.environ['TEST_SERVER'] = TEST_DB_CONFIG['server']
    os.environ['TEST_DATABASE'] = TEST_DB_CONFIG['database']
    os.environ['TEST_USERNAME'] = TEST_DB_CONFIG['username']
    os.environ['TEST_PASSWORD'] = TEST_DB_CONFIG['password']

    yield

    # 3. 清理测试数据库
    conn = pymssql.connect(
        server=TEST_DB_CONFIG['server'],
        database='master',
        user=TEST_DB_CONFIG['username'],
        password=TEST_DB_CONFIG['password']
    )
    cursor = conn.cursor()
    cursor.execute(f"ALTER DATABASE {TEST_DB_CONFIG['database']} SET SINGLE_USER WITH ROLLBACK IMMEDIATE")
    cursor.execute(f"DROP DATABASE {TEST_DB_CONFIG['database']}")
    conn.commit()
    conn.close()

    # 清理环境变量
    for key in ['TEST_SERVER', 'TEST_DATABASE', 'TEST_USERNAME', 'TEST_PASSWORD']:
        del os.environ[key]

def test_integration_db_connection(test_db):
    # 验证真实连接与查询功能
    with db_connection('TEST') as conn:
        cursor = conn.cursor()
        cursor.execute("SELECT 1")
        result = cursor.fetchone()
        assert result[0] == 1

四、日志应用上下文问题解决 #

当函数中使用app.logger.info时，测试需显式激活Flask应用上下文，两种常用方案：

方法1：直接使用app.app_context()上下文管理器 #

from your_flask_app import app

@patch('pymssql.connect')
def test_db_connection_with_log(mock_connect):
    mock_conn = MagicMock()
    mock_connect.return_value = mock_conn

    # 激活应用上下文
    with app.app_context():
        with db_connection('TEST') as conn:
            # 执行测试逻辑
            pass

方法2：用pytest-flask的fixture简化配置 #

先安装依赖：

pip install pytest-flask

测试代码中使用内置fixture自动激活上下文：

import pytest
from your_flask_app import app

@pytest.fixture
def app_context():
    with app.app_context():
        yield

def test_db_connection_with_log(app_context, mock_connect):
    mock_conn = MagicMock()
    mock_connect.return_value = mock_conn
    
    with db_connection('TEST') as conn:
        # 执行测试逻辑
        pass

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