使用moto 5.x.x对操作DynamoDB的AWS Lambda函数进行单元测试

在新版 moto 中，@mock_dynamodb 和 @mock_dynamodb2 已被移除，统一使用 @mock_aws 来模拟所有 AWS 服务。下面是完整的单元测试示例，涵盖 Lambda 与 DynamoDB 的常见交互场景。

前提依赖 #

确保安装必要的包：

pip install moto boto3 pytest

示例 Lambda 函数 #

假设我们有一个处理用户 CRUD 操作的 Lambda，代码文件为 lambda_function.py：

import boto3
import os
from botocore.exceptions import ClientError

# 初始化 DynamoDB 资源
dynamodb = boto3.resource('dynamodb')
TABLE_NAME = os.environ.get('TABLE_NAME', 'Users')
table = dynamodb.Table(TABLE_NAME)

def lambda_handler(event, context):
    action = event.get('action')
    user_id = event.get('user_id')

    try:
        if action == 'create':
            name = event['name']
            table.put_item(Item={'user_id': user_id, 'name': name})
            return {'statusCode': 200, 'body': f"用户 {user_id} 创建成功"}
        
        elif action == 'get':
            response = table.get_item(Key={'user_id': user_id})
            if 'Item' in response:
                return {'statusCode': 200, 'body': response['Item']}
            else:
                return {'statusCode': 404, 'body': "用户不存在"}
        
        elif action == 'update':
            new_name = event['new_name']
            table.update_item(
                Key={'user_id': user_id},
                UpdateExpression='SET #n = :name',
                ExpressionAttributeNames={'#n': 'name'},
                ExpressionAttributeValues={':name': new_name}
            )
            return {'statusCode': 200, 'body': f"用户 {user_id} 信息更新成功"}
        
        else:
            return {'statusCode': 400, 'body': "无效操作类型"}
    
    except ClientError as e:
        return {'statusCode': 500, 'body': f"操作失败: {str(e)}"}

单元测试代码 #

创建测试文件 test_lambda.py，使用 @mock_aws 模拟 AWS 服务，并验证 Lambda 的每一项功能：

import pytest
from lambda_function import lambda_handler
import boto3
from moto import mock_aws
import os

@mock_aws
def test_lambda_dynamodb_full_flow():
    # 1. 配置环境变量（可选，若 Lambda 依赖环境变量）
    os.environ['TABLE_NAME'] = 'Users'

    # 2. 创建模拟的 DynamoDB 表
    dynamodb = boto3.resource('dynamodb', region_name='us-east-1')
    table = dynamodb.create_table(
        TableName='Users',
        KeySchema=[{'AttributeName': 'user_id', 'KeyType': 'HASH'}],
        AttributeDefinitions=[{'AttributeName': 'user_id', 'AttributeType': 'S'}],
        ProvisionedThroughput={'ReadCapacityUnits': 5, 'WriteCapacityUnits': 5}
    )
    table.wait_until_exists()  # 确保表创建完成

    # 3. 测试创建用户
    create_event = {'action': 'create', 'user_id': 'u1001', 'name': '张三'}
    create_resp = lambda_handler(create_event, None)
    assert create_resp['statusCode'] == 200
    assert create_resp['body'] == "用户 u1001 创建成功"

    # 验证表中存在该用户
    db_item = table.get_item(Key={'user_id': 'u1001'})['Item']
    assert db_item['name'] == '张三'

    # 4. 测试获取用户
    get_event = {'action': 'get', 'user_id': 'u1001'}
    get_resp = lambda_handler(get_event, None)
    assert get_resp['statusCode'] == 200
    assert get_resp['body']['user_id'] == 'u1001'
    assert get_resp['body']['name'] == '张三'

    # 5. 测试更新用户
    update_event = {'action': 'update', 'user_id': 'u1001', 'new_name': '张三丰'}
    update_resp = lambda_handler(update_event, None)
    assert update_resp['statusCode'] == 200
    assert update_resp['body'] == "用户 u1001 信息更新成功"

    # 验证更新结果
    updated_item = table.get_item(Key={'user_id': 'u1001'})['Item']
    assert updated_item['name'] == '张三丰'

    # 6. 测试获取不存在的用户
    get_nonexistent_event = {'action': 'get', 'user_id': 'u9999'}
    get_nonexistent_resp = lambda_handler(get_nonexistent_event, None)
    assert get_nonexistent_resp['statusCode'] == 404
    assert get_nonexistent_resp['body'] == "用户不存在"

    # 7. 测试无效操作
    invalid_event = {'action': 'delete'}
    invalid_resp = lambda_handler(invalid_event, None)
    assert invalid_resp['statusCode'] == 400
    assert invalid_resp['body'] == "无效操作类型"

关键注意事项 #

• 模拟环境初始化：每次测试都会启动一个全新的模拟 AWS 环境，因此需要在测试中预先创建 Lambda 依赖的 DynamoDB 表。
• 区域一致性：确保测试中使用的 AWS 区域与 Lambda 代码中的区域一致（若未指定，moto 默认使用 us-east-1）。
• 状态验证：可以直接使用 boto3 客户端/资源在测试中查询模拟的 DynamoDB，验证 Lambda 操作的结果是否符合预期。
• 异常处理：测试中可以模拟 DynamoDB 异常（如权限错误），验证 Lambda 的错误处理逻辑。

使用 unittest 的替代写法 #

如果你习惯使用 Python 标准库的 unittest，可以这样编写测试：

import unittest
from lambda_function import lambda_handler
import boto3
from moto import mock_aws
import os

class TestLambdaDynamoDB(unittest.TestCase):
    @mock_aws
    def test_create_and_update_user(self):
        # 创建表
        dynamodb = boto3.resource('dynamodb')
        dynamodb.create_table(
            TableName='Users',
            KeySchema=[{'AttributeName': 'user_id', 'KeyType': 'HASH'}],
            AttributeDefinitions=[{'AttributeName': 'user_id', 'AttributeType': 'S'}],
            ProvisionedThroughput={'ReadCapacityUnits': 5, 'WriteCapacityUnits': 5}
        )

        # 测试创建
        create_event = {'action': 'create', 'user_id': 'u2001', 'name': '李四'}
        create_resp = lambda_handler(create_event, None)
        self.assertEqual(create_resp['statusCode'], 200)

        # 测试更新
        update_event = {'action': 'update', 'user_id': 'u2001', 'new_name': '李四光'}
        update_resp = lambda_handler(update_event, None)
        self.assertEqual(update_resp['statusCode'], 200)

        # 验证更新结果
        table = dynamodb.Table('Users')
        updated_item = table.get_item(Key={'user_id': 'u2001'})['Item']
        self.assertEqual(updated_item['name'], '李四光')

if __name__ == '__main__':
    unittest.main()

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