Skip to content

Instantly share code, notes, and snippets.

@miketheman
Created February 28, 2023 21:51
Show Gist options
  • Save miketheman/90f482df32818f45979817738f8d2cd7 to your computer and use it in GitHub Desktop.
Save miketheman/90f482df32818f45979817738f8d2cd7 to your computer and use it in GitHub Desktop.
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