How to mock boto3 resource method in a while loop for unit testing

Question

I'm trying to unit test some code that I wrote with the boto3 library and cannot seem to figure out how to do it.

This is the code I'm trying to test

def handler(event, context):
    instance_id = event['instance_id']
    ec2 = boto3.resource('ec2')
    instance = ec2.Instance(instance_id)

    instance.start()
    while instance.state['Name'] != 'running':
        sleep(1)
        instance.reload()

And here's how the test looks currently.

@patch.object(module.boto3, "resource")
def test_starts_instance(resource_mock):
    event = {"instance_id": "id-1234567"}

    instance_mock = resource_mock.return_value.Instance.return_value
    instance_mock.state = [{"Name": "stopping"}]

    handler(event, None)

    resource_mock.assert_called_once_with("ec2")
    instance_mock.start.assert_called_once()
    instance_mock.reload.assert_called_once()

I just want to change the instance.state['Name'] to 'running' when instance.reload() is called so that the test actually tests the sleep and reload calls.

Answer 1

Here is what you can do:

@patch.object(module.boto3, "resource")
def test_starts_instance(resource_mock):
    event = {"instance_id": "id-1234567"}

    instance_mock = resource_mock.return_value.Instance.return_value
    instance_mock.state = {"Name": "stopping"}

    def reload():
        instance_mock.state = {"Name": "running"}

    instance_mock.reload.side_effect = reload

    handler(event, None)

    resource_mock.assert_called_once_with("ec2")
    instance_mock.start.assert_called_once()
    instance_mock.reload.assert_called_once()

You replace the reload function of the instance with your implementation, which sets the status to be the desired one so the loop can stop.