使用pytest mock Secrets Manager测试Aurora密码轮换Lambda

你原有测试代码存在以下语法和逻辑错误：

• 模块导入语法错误：导入自定义模块不需要加.py后缀，import lambda_function.py as testpass应改为import lambda_function as testpass
• moto导入语法错误：import moto import mock_secretsmanager应改为from moto import mock_secretsmanager
• 参数传值错误：service_client入参需要传入初始化完成的boto3 Secrets Manager客户端实例，不能直接传字符串"secretsmanager"
• 依赖未模拟：既没有在mock的Secrets Manager中预置对应AWSPENDING阶段的测试密钥，也没有mock数据库连接逻辑，运行时既找不到对应密钥，也会尝试真实连接数据库导致失败
• 变量重复定义：重复声明了token变量，后定义的值会覆盖前值

测试实现思路 #

test_secret只依赖两个外部组件，全部通过mock隔离即可实现离线单元测试：

1. Secrets Manager服务：用moto库模拟，无需真实调用AWS接口，在mock环境中预置符合要求的测试密钥即可
2. 数据库连接逻辑：通过pytest-mock直接patch代码中的get_connection函数，返回模拟的数据库连接对象，验证SQL执行、连接关闭等逻辑是否符合预期，不需要真实连接数据库

前置依赖安装 #

先安装测试需要的包：

pip install pytest moto boto3 pytest-mock

完整测试代码 #

import boto3
import json
import pytest
from moto import mock_secretsmanager
# 导入被测Lambda模块，注意不要加.py后缀
import lambda_function as rotation_lambda

# 测试用固定常量，全部为构造的假值，无需和真实环境一致
TEST_REGION = "us-east-1"
TEST_SECRET_ARN = "arn:aws:secretsmanager:us-east-1:123456789012:secret:test-aurora-secret-Xyz123"
TEST_TOKEN = "test-client-token-abcdef123456"
# 符合轮换Lambda要求的测试密钥结构
TEST_PENDING_CREDENTIALS = {
    "host": "test-aurora-cluster.cluster-xxx.us-east-1.rds.amazonaws.com",
    "port": 3306,
    "username": "db_admin",
    "password": "test_rotated_password_456",
    "dbname": "app_db"
}

@mock_secretsmanager
def test_test_secret_when_connection_succeed(mocker):
    # 1. 初始化mock的Secrets Manager客户端，预置AWSPENDING阶段的测试密钥
    sm_client = boto3.client("secretsmanager", region_name=TEST_REGION)
    sm_client.create_secret(Name="test-aurora-secret", ARN=TEST_SECRET_ARN)
    sm_client.put_secret_value(
        SecretId=TEST_SECRET_ARN,
        ClientRequestToken=TEST_TOKEN,
        SecretString=json.dumps(TEST_PENDING_CREDENTIALS),
        VersionStages=["AWSPENDING"]
    )

    # 2. 构造mock数据库连接对象，patch掉真实的get_connection方法
    mock_db_conn = mocker.MagicMock()
    mock_cursor = mocker.MagicMock()
    # 模拟with conn.cursor()的上下文管理器行为
    mock_db_conn.cursor.return_value.__enter__.return_value = mock_cursor
    mocker.patch.object(rotation_lambda, "get_connection", return_value=mock_db_conn)

    # 3. 调用被测函数
    rotation_lambda.test_secret(sm_client, TEST_SECRET_ARN, TEST_TOKEN)

    # 4. 断言所有逻辑符合预期
    # 验证get_connection接收到了正确的密钥参数
    rotation_lambda.get_connection.assert_called_once_with(TEST_PENDING_CREDENTIALS)
    # 验证执行了权限检查SQL
    mock_cursor.execute.assert_called_once_with("SELECT NOW()")
    mock_db_conn.commit.assert_called_once()
    # 验证连接最终被关闭
    mock_db_conn.close.assert_called_once()

@mock_secretsmanager
def test_test_secret_when_connection_failed(mocker):
    # 初始化mock SM环境，预置测试密钥
    sm_client = boto3.client("secretsmanager", region_name=TEST_REGION)
    sm_client.create_secret(Name="test-aurora-secret", ARN=TEST_SECRET_ARN)
    sm_client.put_secret_value(
        SecretId=TEST_SECRET_ARN,
        ClientRequestToken=TEST_TOKEN,
        SecretString=json.dumps(TEST_PENDING_CREDENTIALS),
        VersionStages=["AWSPENDING"]
    )

    # 模拟get_connection返回None，即数据库连接失败
    mocker.patch.object(rotation_lambda, "get_connection", return_value=None)

    # 验证函数抛出预期的ValueError
    with pytest.raises(ValueError, match=f"Unable to log into database with pending secret of secret ARN {TEST_SECRET_ARN}"):
        rotation_lambda.test_secret(sm_client, TEST_SECRET_ARN, TEST_TOKEN)

扩展测试场景 #

如果需要覆盖其他异常分支，只需要调整mock逻辑即可：

• 测试密钥缺少必填字段：在put_secret_value时传入缺少host/username等字段的SecretString，断言抛出KeyError
• 测试密钥不是合法JSON：传入非JSON格式的SecretString，断言抛出ValueError
• 测试指定ARN/阶段的密钥不存在：不要预置对应VersionStages为AWSPENDING的密钥版本，断言抛出对应资源不存在的异常

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