Created
February 28, 2023 21:51
-
-
Save miketheman/90f482df32818f45979817738f8d2cd7 to your computer and use it in GitHub Desktop.
Mock side_effect https://docs.python.org/3/library/unittest.mock.html#unittest.mock.Mock.side_effect
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
from unittest.mock import patch | |
import boto3 | |
from botocore.stub import Stubber | |
# create an ssm client | |
ssm = boto3.client("ssm") | |
# create an sqs client | |
sqs = boto3.client("sqs") | |
# create a Stubber response for ssm | |
ssm_stubber = Stubber(ssm) | |
ssm_stubber.add_response( | |
"get_parameter", | |
{ | |
"Parameter": { | |
"Name": "my_parameter", | |
"Type": "String", | |
"Value": "my_value", | |
"Version": 1, | |
"LastModifiedDate": 1234567890.123, | |
"ARN": "arn:aws:ssm:us-east-1:123456789012:parameter/my_parameter", | |
} | |
}, | |
{"Name": "my_parameter", "WithDecryption": True}, | |
) | |
ssm_stubber.activate() | |
# create a Stubber response for sqs | |
sqs_stubber = Stubber(sqs) | |
sqs_stubber.add_response( | |
"get_queue_url", | |
{"QueueUrl": "https://sqs.us-east-1.amazonaws.com/123456789012/my_queue"}, | |
{"QueueName": "my_queue", "QueueOwnerAWSAccountId": "123456789012"}, | |
) | |
sqs_stubber.activate() | |
# patch boto3.client with a side_effect that switches based on the client name | |
with patch("boto3.client") as mock_client: | |
mock_client.side_effect = lambda svc_name: { | |
"ssm": ssm_stubber.client, | |
"sqs": sqs_stubber.client, | |
}.get(svc_name, boto3.client(svc_name)) | |
# run the code that uses the ssm and sqs clients | |
ssm_param = ssm.get_parameter(Name="my_parameter", WithDecryption=True) | |
sqs_queue = sqs.get_queue_url( | |
QueueName="my_queue", QueueOwnerAWSAccountId="123456789012" | |
) | |
# assert that the stubbed responses were called | |
ssm_stubber.assert_no_pending_responses() | |
sqs_stubber.assert_no_pending_responses() | |
assert ssm_param["Parameter"]["Value"] == "my_value" | |
assert ( | |
sqs_queue["QueueUrl"] == "https://sqs.us-east-1.amazonaws.com/123456789012/my_queue" | |
) |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment