如何用Python Unittest覆盖DynamoDB分页扫描代码分支？

原函数代码 #

import boto3

# 假设TABLE是预先初始化的DynamoDB表对象
TABLE = boto3.resource('dynamodb').Table('employers')

def get_employer_list():
    employer_list = []
    response = TABLE.scan()
    employer_list.extend(response['Items'])
    
    while 'LastEvaluatedKey' in response:
        # 需要覆盖的两行代码
        response = TABLE.scan(ExclusiveStartKey=response['LastEvaluatedKey'])
        employer_list.extend(response['Items'])
    
    return employer_list

现有测试代码（未覆盖分页分支） #

import unittest
from unittest.mock import patch
from your_module import get_employer_list, TABLE

class TestEmployerList(unittest.TestCase):
    @patch.object(TABLE, 'scan')
    def test_get_employer_list_single_page(self, mock_scan):
        # 模拟单次返回所有数据，无LastEvaluatedKey
        mock_scan.return_value = {
            'Items': [{'id': '1', 'name': 'Company A'}, {'id': '2', 'name': 'Company B'}],
        }
        
        result = get_employer_list()
        
        self.assertEqual(len(result), 2)
        mock_scan.assert_called_once()

完善后的测试代码（覆盖分页分支） #

要触发分页逻辑，需要让TABLE.scan被调用两次：第一次返回带LastEvaluatedKey的响应，第二次返回剩余数据且无该键。通过side_effect可以模拟多次调用的不同返回值：

import unittest
from unittest.mock import patch
from your_module import get_employer_list, TABLE

class TestEmployerList(unittest.TestCase):
    @patch.object(TABLE, 'scan')
    def test_get_employer_list_with_pagination(self, mock_scan):
        # 模拟第一次scan返回部分数据和分页标记
        first_response = {
            'Items': [{'id': '1', 'name': 'Company A'}, {'id': '2', 'name': 'Company B'}],
            'LastEvaluatedKey': {'id': '2'}
        }
        # 模拟第二次scan返回剩余数据，无分页标记
        second_response = {
            'Items': [{'id': '3', 'name': 'Company C'}, {'id': '4', 'name': 'Company D'}]
        }
        
        # 设置mock的多次返回值
        mock_scan.side_effect = [first_response, second_response]
        
        result = get_employer_list()
        
        # 验证结果合并了两次分页的数据
        self.assertEqual(len(result), 4)
        self.assertEqual(result[2]['name'], 'Company C')
        
        # 验证scan被调用了两次
        self.assertEqual(mock_scan.call_count, 2)
        
        # 验证第二次调用正确传入了分页参数
        mock_scan.assert_any_call()
        mock_scan.assert_any_call(ExclusiveStartKey={'id': '2'})

关键说明 #

• side_effect可以让mock对象在多次调用时返回不同结果，完美模拟DynamoDB的分页返回逻辑
• 通过call_count确认分页逻辑被触发，确保目标代码行被执行
• assert_any_call验证第二次调用时正确传入了ExclusiveStartKey，覆盖需要测试的分支
• 结果验证确保两次分页的数据被正确合并

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