Python单元测试中用mock patch修补sys.argv遇IndexError的解决方法

问题根源 #

你的seed_dynamodb.py在模块顶层直接执行sys.argv[1]和sys.argv[2]的读取操作，当测试模块导入seed_dynamodb时，这部分代码会立即执行——此时sys.argv还未被patch，自然触发IndexError。

第一步：修复原脚本的执行逻辑 #

将命令行参数读取逻辑移到if __name__ == "__main__":代码块中，这样只有直接运行脚本时才会执行参数读取，导入模块时不会触发：

import sys
import boto3

def main(env, region):
    dynamodb_client = boto3.client('dynamodb')
    timestamp = '1234567'
    table_name = 'syn-fcad-nielsen-' + env + '-time'
    print(f'{table_name=}')
    
    if env == 'uat':
        timestamp = 1234567
    if env == 'prod':
        timestamp = 1234567
        
    response = dynamodb_client.get_item(
        TableName=table_name, 
        Key={'BaseTime': {'S': 'Timestamp'}}
    )
    
    if 'Item' in response:
        print("Item exists in Dynamo DB table")
        # 修正原代码笔误：response['Items'] -> response['Item']
        timestamp = response['Item']['Value']['N']
    else: 
        response = dynamodb_client.put_item(
            TableName=table_name, 
            Item={
                'BaseTime': {'S': 'Timestamp'},
                # DynamoDB的N类型需要字符串，确保类型正确
                'Value': {'N': str(timestamp)}
            }
        )

# 仅在直接运行脚本时读取命令行参数
if __name__ == "__main__":
    env = sys.argv[1]
    region = sys.argv[2]
    l = len(sys.argv)
    print(f'{env=}{region=}{l=}')
    main(env, region)

第二步：编写正确的单元测试 #

测试时直接调用main函数并传入参数，比patchsys.argv更简洁可控。如果需要测试命令行参数解析逻辑，可单独编写测试用例：

import unittest
from unittest.mock import patch, Mock
import sys
from module import seed_dynamodb

class SeedDynamoDBTest(unittest.TestCase):
    @patch('module.seed_dynamodb.boto3.client')
    def test_item_exists_scenario(self, mock_boto_client):
        # 配置mock DynamoDB客户端
        mock_dynamodb = Mock()
        mock_boto_client.return_value = mock_dynamodb
        
        # 模拟Item已存在的返回结果
        mock_dynamodb.get_item.return_value = {
            "Item": {"Value": {"N": "1678230539"}}
        }
        
        # 直接调用main函数，传入测试参数
        seed_dynamodb.main("uat", "us-east-1")
        
        # 验证方法调用是否符合预期
        mock_dynamodb.get_item.assert_called_once_with(
            TableName='syn-fcad-nielsen-uat-time',
            Key={'BaseTime': {'S': 'Timestamp'}}
        )
        mock_dynamodb.put_item.assert_not_called()

    @patch('module.seed_dynamodb.boto3.client')
    def test_item_not_exists_scenario(self, mock_boto_client):
        mock_dynamodb = Mock()
        mock_boto_client.return_value = mock_dynamodb
        
        # 模拟Item不存在的返回结果
        mock_dynamodb.get_item.return_value = {}
        
        seed_dynamodb.main("prod", "us-east-1")
        
        # 验证put_item是否被正确调用
        mock_dynamodb.put_item.assert_called_once_with(
            TableName='syn-fcad-nielsen-prod-time',
            Item={
                'BaseTime': {'S': 'Timestamp'},
                'Value': {'N': '1234567'}
            }
        )
        mock_dynamodb.get_item.assert_called_once()

    # 单独测试命令行参数解析逻辑
    @patch('sys.argv', ['seed_dynamodb.py', 'pr', 'us-west-2'])
    @patch('module.seed_dynamodb.main')
    def test_command_line_args_parsing(self, mock_main):
        # 重新加载模块，触发main guard中的参数读取逻辑
        import importlib
        importlib.reload(seed_dynamodb)
        
        # 验证main函数是否收到正确的参数
        mock_main.assert_called_once_with('pr', 'us-west-2')

if __name__ == "__main__":
    unittest.main(verbosity=2)

关键说明 #

1. 模块顶层代码仅在导入时执行一次，移到main guard中可彻底避免导入时的副作用。
2. 直接调用main函数传参，比patchsys.argv更符合单元测试的单一职责原则，测试逻辑更清晰。
3. 修正了原代码中response['Items']的笔误，以及DynamoDBN类型要求字符串的细节问题。

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