How to use botocore.stub.Stubber to stub out the response of an SFN client without affecting the behavior of an SQS client

Question

I am using stubber to stub out the response of "get_execution_history" method of stepfunction but it looks like it is affecting sqs client too. I have used moto library to mock 'sqs' service.

I have my lambda_function.py

stepfunctionclient = boto3.client('stepfunctions', region_name='us-east-1')
sqsClient = boto3.client('sqs', region_name='us-east-1')

def lambda_handler(event, context):
    # calling get_execution_history method
    responseType = stepfunctionclient.get_execution_history(
                   executionArn="my_step_arn",
                   maxResults=1000,
                   reverseOrder=True,
                 )
    # sending messages to sqs queue
    sqsClient.send_message(QueueUrl="my_url", MessageBody="some_message")

then I have my test_lambdaHandler.py

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

    sfClient = boto3.client("stepfunctions", region_name="us-east-1")
    stubber = botocore.stub.Stubber(sfClient)
    execution_history = {
           # mock response here
    }

    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):
        with mock.patch(boto3,'client', mock.MagicMock(return_value=sfClient)):
             module = importlib.import_module(
                  "module_name"
                  ).lambda_function

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

The error I am getting is

AttributeError: 'SFN' object has no attribute 'send_message'

Can anyone please help me here to stub out the response of get_execution_history method without affecting the sqs client.

Answer 1

You are mocking the boto3.client function to return the sfClient for your SQS client. Just create a custom side effect function for the boto3.client mock that returns the correct client based on the service_name argument.

from unittest.mock import MagicMock

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

    sfClient = boto3.client("stepfunctions", region_name="us-east-1")
    sqsClient = boto3.client("sqs", region_name="us-east-1")

    stubber = botocore.stub.Stubber(sfClient)
    execution_history = {
        # mock response here
    }

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

    stubber.activate()

    def client_side_effect(service_name, *args, **kwargs):
        if service_name == "stepfunctions":
            return sfClient
        elif service_name == "sqs":
            return sqsClient
        else:
            return boto3.client(service_name, *args, **kwargs)

    with mock.patch.dict(os.environ, dummyenvironmentVariables, clear=True):
        with mock.patch("boto3.client", side_effect=client_side_effect):
            module = importlib.import_module("module_name").lambda_function
            actualReturn = module.lambda_handler(event, None)
            print(actualReturn)