如何使用botocore.stub.Stubber stub SFN客户端且不影响SQS客户端

你用mock.patch(boto3, 'client', mock.MagicMock(return_value=sfClient))全局替换了boto3.client的实现，导致Lambda代码中创建SQS客户端时，拿到的是被Stubber处理的StepFunctions客户端，而StepFunctions客户端本身没有send_message方法，因此触发报错。

以下两种方法都能实现仅Stub StepFunctions的get_execution_history方法，同时保留moto对SQS的Mock效果：

方法一：按服务名返回对应客户端 #

修改Mock逻辑，让boto3.client根据传入的服务参数返回不同客户端：

@mock_sqs
def test_lambdaHandler(
    event,
    sqsQueue,
    dummyenvironmentVariables,
 ):
    sqsQueue()

    # 创建并Stub StepFunctions客户端
    sfClient = boto3.client("stepfunctions", region_name='us-east-1')
    stubber = botocore.stub.Stubber(sfClient)
    execution_history = {
           # 填入你的Mock响应内容
    }

    stubber.add_response(
         'get_execution_history',
          execution_history,
          expected_params={
               'executionArn': "my_step_arn",
               'maxResults':1000,
               'reverseOrder':True
              }
         )
    stubber.activate()

    # 创建moto Mock的SQS客户端
    sqs_client_mock = boto3.client('sqs', region_name='us-east-1')

    def mock_boto_client(service_name, **kwargs):
        if service_name == 'stepfunctions':
            return sfClient
        elif service_name == 'sqs':
            return sqs_client_mock
        # 其他服务可按需处理
        return boto3.client(service_name, **kwargs)

    with mock.patch.dict(os.environ, dummyenvironmentVariables, clear=True):
        with mock.patch('boto3.client', side_effect=mock_boto_client):
             module = importlib.import_module(
                  "module_name"
                  ).lambda_function

             actualReturn = module.lambda_handler(event, None)
             print(actualReturn)

方法二：直接Patch Lambda模块的StepFunctions客户端变量 #

针对你的代码中客户端是模块级全局变量的场景，直接替换Lambda模块里的stepfunctionclient实例：

@mock_sqs
def test_lambdaHandler(
    event,
    sqsQueue,
    dummyenvironmentVariables,
 ):
    sqsQueue()

    # 创建并Stub StepFunctions客户端
    sfClient = boto3.client("stepfunctions", region_name='us-east-1')
    stubber = botocore.stub.Stubber(sfClient)
    execution_history = {
           # 填入你的Mock响应内容
    }

    stubber.add_response(
         'get_execution_history',
          execution_history,
          expected_params={
               'executionArn': "my_step_arn",
               'maxResults':1000,
               'reverseOrder':True
              }
         )
    stubber.activate()

    with mock.patch.dict(os.environ, dummyenvironmentVariables, clear=True):
        lambda_module = importlib.import_module("module_name").lambda_function
        # 直接替换Lambda模块中的全局StepFunctions客户端
        with mock.patch.object(lambda_module, 'stepfunctionclient', sfClient):
             actualReturn = lambda_module.lambda_handler(event, None)
             print(actualReturn)

• 方法一适合动态创建客户端的场景，方法二更适合客户端为模块级全局变量的情况（你的代码属于此类）。
• 两种方法都避免了全局替换boto3.client，从而保留了moto对SQS的Mock效果。

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