def mock_s3_get_object(bucket_name, granted_prefixes, key, success_response):
_keep_it_real()
client = boto3.real_client("s3")
stubber = Stubber(client)
if any([key.startswith(prefix) for prefix in granted_prefixes]):
stubber.add_response(
"get_object",
success_response,
{
"Bucket": bucket_name,
"Key": key
},
)
else:
stubber.add_client_error("get_object",
expected_params={
"Bucket": bucket_name,
"Key": key
})
# replace the get_presigned_url so it runs without AWS creds
client.generate_presigned_url = lambda op, Params, ExpiresIn, HttpMethod: fake_url(
Params["Bucket"], Params["Key"])
stubber.activate()
# override boto.client to return the mock client
boto3.client = lambda service, region_name=None, config=None: client
return stubber
def mock_cognito_create_user_add_user_to_group_fails(admin_user,
create_user_arguments):
_keep_it_real()
client = boto3.real_client("cognito-idp")
group_name = admin_user["group"]["value"]
stubber = Stubber(client)
# Add responses
stub_response_cognito_admin_create_user(stubber, admin_user,
create_user_arguments)
stub_response_cognito_admin_set_user_mfa_preference(
stubber, admin_user["email"])
stub_response_cognito_admin_set_user_settings(stubber, admin_user["email"])
# stub_response_cognito_admin_add_user_to_group(
# stubber, admin_user["email"], group_name
# )
stubber.add_client_error(
"admin_add_user_to_group",
expected_params={
"UserPoolId": MOCK_COGNITO_USER_POOL_ID,
"Username": admin_user["email"],
"GroupName": group_name,
},
)
stub_response_cognito_admin_disable_user(stubber, admin_user["email"])
stubber.activate()
# override boto.client to return the mock client
boto3.client = lambda service, region_name=None: client
return stubber
def mock_lambda(mock_get_function_response, mock_list_tags_response):
""" Mock lambda get function and list tags """
_keep_it_real()
region = "eu-west-2"
client = boto3.real_client("lambda", region_name=region)
stubber = Stubber(client)
# mock get_function response
get_function_params = {"FunctionName": "lambda-function"}
stubber.add_response(
"get_function", mock_get_function_response, get_function_params
)
list_tags_params = {
"Resource": "arn:aws:lambda:eu-west-2:123456789012:function:lambda-function"
}
stubber.add_response("list_tags", mock_list_tags_response, list_tags_params)
stubber.activate()
# override boto.client to return the mock client
boto3.client = lambda service, region_name: client
return stubber
def mock_ssm_get_parameters_by_path():
_keep_it_real()
client = boto3.real_client("ssm")
stubber = Stubber(client)
# mock get_function response
mock_get_parameters_by_path_page_1 = {
"Parameters": [
{
"Name": "/alert-processor/tokens/github/user-a",
"Value": "github-token-a",
},
{
"Name": "/alert-processor/tokens/github/user-b",
"Value": "github-token-b",
},
],
"NextToken":
"page-2",
}
stubber.add_response(
"get_parameters_by_path",
mock_get_parameters_by_path_page_1,
{
"Path": "/alert-processor/tokens/github/",
"Recursive": True,
"WithDecryption": True,
},
)
mock_get_parameters_by_path_page_2 = {
"Parameters": [
{
"Name": "/alert-processor/tokens/github/user-c",
"Value": "github-token-c",
},
{
"Name": "/alert-processor/tokens/github/user-d",
"Value": "github-token-d",
},
]
}
stubber.add_response(
"get_parameters_by_path",
mock_get_parameters_by_path_page_2,
{
"Path": "/alert-processor/tokens/github/",
"Recursive": True,
"WithDecryption": True,
"NextToken": "page-2",
},
)
stubber.activate()
# override boto.client to return the mock client
boto3.client = lambda service: client
return stubber
def mock_cognito_create_user_set_mfa_fails(admin_user, create_user_arguments):
_keep_it_real()
client = boto3.real_client("cognito-idp")
group_name = admin_user["group"]["value"]
stubber = Stubber(client)
# Add responses
stub_response_cognito_list_user_pools(stubber)
stub_response_cognito_admin_create_user(stubber, admin_user, create_user_arguments)
stub_response_cognito_list_user_pools(stubber)
# Replace admin_set_user_mfa_preference response with ClientError
# stub_response_cognito_admin_set_user_mfa_preference(stubber, admin_user["email"])
stubber.add_client_error(
"admin_set_user_mfa_preference",
expected_params={
"SMSMfaSettings": {"Enabled": True, "PreferredMfa": True},
"UserPoolId": MOCK_COGNITO_USER_POOL_ID,
"Username": admin_user["email"],
},
)
stub_response_cognito_list_user_pools(stubber)
stub_response_cognito_admin_set_user_settings(stubber, admin_user["email"])
stub_response_cognito_list_user_pools(stubber)
stub_response_cognito_admin_add_user_to_group(
stubber, admin_user["email"], group_name
)
stub_response_cognito_list_user_pools(stubber)
stub_response_cognito_admin_disable_user(stubber, admin_user["email"])
stubber.activate()
# override boto.client to return the mock client
boto3.client = lambda service, region_name=None: client
return stubber
def mock_delete_user_failure(admin_user, admin_get_user):
_keep_it_real()
client = boto3.real_client("cognito-idp")
stubber = Stubber(client)
# Add responses
# get user
stub_response_cognito_list_user_pools(stubber)
stub_response_cognito_admin_get_user(stubber, admin_user["email"], admin_get_user)
stub_response_cognito_list_user_pools(stubber)
stub_response_cognito_admin_list_groups_for_user(stubber, admin_user["email"])
# delete user
stub_response_cognito_list_user_pools(stubber)
stubber.add_client_error(
"admin_delete_user",
expected_params={
"UserPoolId": MOCK_COGNITO_USER_POOL_ID,
"Username": admin_user["email"],
},
)
stubber.activate()
# override boto.client to return the mock client
boto3.client = lambda service, region_name=None: client
return stubber
def mock_cognito_create_user_set_user_settings_fails(admin_user, create_user_arguments):
_keep_it_real()
client = boto3.real_client("cognito-idp")
group_name = admin_user["group"]["value"]
stubber = Stubber(client)
# Add responses
stub_response_cognito_list_user_pools(stubber)
stub_response_cognito_admin_create_user(stubber, admin_user, create_user_arguments)
stub_response_cognito_list_user_pools(stubber)
stub_response_cognito_admin_set_user_mfa_preference(stubber, admin_user["email"])
stub_response_cognito_list_user_pools(stubber)
# stub_response_cognito_admin_set_user_settings(stubber, admin_user["email"])
stubber.add_client_error(
"admin_set_user_settings",
expected_params={
"MFAOptions": [{"DeliveryMedium": "SMS", "AttributeName": "phone_number"}],
"UserPoolId": MOCK_COGNITO_USER_POOL_ID,
"Username": admin_user["email"],
},
)
stub_response_cognito_list_user_pools(stubber)
stub_response_cognito_admin_add_user_to_group(
stubber, admin_user["email"], group_name
)
stub_response_cognito_list_user_pools(stubber)
stub_response_cognito_admin_disable_user(stubber, admin_user["email"])
stubber.activate()
# override boto.client to return the mock client
boto3.client = lambda service, region_name=None: client
return stubber
def mock_user_reinvite(admin_user, admin_get_user, create_user_arguments):
_keep_it_real()
client = boto3.real_client("cognito-idp")
group_name = admin_user["group"]["value"]
stubber = Stubber(client)
# Add responses
# get user
stub_response_cognito_list_user_pools(stubber)
stub_response_cognito_admin_get_user(stubber, admin_user["email"], admin_get_user)
stub_response_cognito_list_user_pools(stubber)
stub_response_cognito_admin_list_groups_for_user(stubber, admin_user["email"])
# delete user
stub_response_cognito_list_user_pools(stubber)
stub_response_cognito_admin_delete_user(stubber, admin_user["email"])
# create user
stub_response_cognito_list_user_pools(stubber)
stub_response_cognito_admin_create_user(stubber, admin_user, create_user_arguments)
stub_response_cognito_list_user_pools(stubber)
stub_response_cognito_admin_set_user_mfa_preference(stubber, admin_user["email"])
stub_response_cognito_list_user_pools(stubber)
stub_response_cognito_admin_set_user_settings(stubber, admin_user["email"])
stub_response_cognito_list_user_pools(stubber)
stub_response_cognito_admin_add_user_to_group(
stubber, admin_user["email"], group_name
)
stubber.activate()
# override boto.client to return the mock client
boto3.client = lambda service, region_name=None: client
return stubber
def mock_create_user_failure(admin_user, admin_get_user, create_user_arguments):
_keep_it_real()
client = boto3.real_client("cognito-idp")
stubber = Stubber(client)
# Add responses
# get user
stub_response_cognito_list_user_pools(stubber)
stub_response_cognito_admin_get_user(stubber, admin_user["email"], admin_get_user)
stub_response_cognito_list_user_pools(stubber)
stub_response_cognito_admin_list_groups_for_user(stubber, admin_user["email"])
# delete user
stub_response_cognito_list_user_pools(stubber)
stub_response_cognito_admin_delete_user(stubber, admin_user["email"])
# create user
stub_response_cognito_list_user_pools(stubber)
stubber.add_client_error(
"admin_create_user", expected_params=create_user_arguments,
)
stubber.activate()
# override boto.client to return the mock client
boto3.client = lambda service, region_name=None: client
return stubber
def mock_cloudwatch(mock_get_metric_statistics):
"""Mock Cloudwatch metric stats"""
_keep_it_real()
region = "eu-west-2"
client = boto3.real_client("cloudwatch", region_name=region)
stubber = Stubber(client)
get_cloudwatch_params = {
"Namespace": "AWS/Lambda",
"MetricName": "Errors",
"Dimensions": [{"Name": "FunctionName", "Value": "lambda-function"}],
"StartTime": ANY,
"EndTime": ANY,
"Period": (60 * 60 * 24 * 28) + 300, # 28 days + 5 mins (see generic_helper.py)
"Statistics": ["Maximum"],
}
stubber.add_response(
"get_metric_statistics", mock_get_metric_statistics, get_cloudwatch_params
)
stubber.activate()
# override boto.client to return the mock client
boto3.client = lambda service, region_name: client
return stubber
def mock_cognito_admin_set_user_mfa_preference(email):
_keep_it_real()
client = boto3.real_client("cognito-idp")
stubber = Stubber(client)
stub_response_cognito_admin_set_user_mfa_preference(stubber, email)
stubber.activate()
# override boto.client to return the mock client
boto3.client = lambda service, region_name=None: client
return stubber
def mock_cognito_list_pools(env="development"):
_keep_it_real()
client = boto3.real_client("cognito-idp")
stubber = Stubber(client)
# Add responses
stub_response_cognito_list_user_pools(stubber, env)
stubber.activate()
# override boto.client to return the mock client
boto3.client = lambda service, region_name=None: client
return stubber
def mock_config_load_ssm_parameters(path, parameters):
_keep_it_real()
client = boto3.real_client("ssm")
stubber = Stubber(client)
# Add responses
stub_response_ssm_get_parameters_by_path(stubber, path, parameters)
stubber.activate()
# override boto.client to return the mock client
boto3.client = lambda service, region_name=None: client
return stubber
def mock_cognito_admin_add_user_to_group(email, group_name):
_keep_it_real()
client = boto3.real_client("cognito-idp")
stubber = Stubber(client)
# Add responses
stub_response_cognito_admin_add_user_to_group(stubber, email, group_name)
stubber.activate()
# override boto.client to return the mock client
boto3.client = lambda service, region_name=None: client
return stubber
def mock_cognito_admin_update_user_attributes(user, attributes):
_keep_it_real()
client = boto3.real_client("cognito-idp")
stubber = Stubber(client)
# Add responses
stub_response_cognito_admin_update_user_attributes(stubber, user["email"],
attributes)
stubber.activate()
# override boto.client to return the mock client
boto3.client = lambda service, region_name=None: client
return stubber
def mock_cognito_admin_list_groups_for_user(admin_user):
_keep_it_real()
client = boto3.real_client("cognito-idp")
stubber = Stubber(client)
# Add responses
stub_response_cognito_list_user_pools(stubber)
stub_response_cognito_admin_list_groups_for_user(stubber, admin_user["email"])
stubber.activate()
# override boto.client to return the mock client
boto3.client = lambda service, region_name=None: client
return stubber
def mock_cognito_auth_flow(token, test_user):
_keep_it_real()
client = boto3.real_client("cognito-idp")
stubber = Stubber(client)
# Add responses
stub_response_cognito_get_user(stubber, token, test_user)
stub_response_cognito_admin_list_groups_for_user(stubber,
test_user["Username"])
stubber.activate()
# override boto.client to return the mock client
boto3.client = lambda service, region_name=None: client
return stubber
def mock_config_set_app_settings(parameters):
_keep_it_real()
client = boto3.real_client("cognito-idp")
stubber = Stubber(client)
# Add responses
stub_response_cognito_list_user_pool_clients(stubber, parameters)
stub_response_cognito_describe_user_pool_client(stubber, parameters)
stub_response_cognito_describe_user_pool(stubber, parameters)
stubber.activate()
# override boto.client to return the mock client
boto3.client = lambda service, region_name=None: client
return stubber
def mock_user_not_found(email):
_keep_it_real()
client = boto3.real_client("cognito-idp")
stubber = Stubber(client)
# Add responses
stub_response_cognito_list_user_pools(stubber)
stubber.add_client_error(
"admin_get_user",
expected_params={"UserPoolId": MOCK_COGNITO_USER_POOL_ID, "Username": email},
)
stubber.activate()
# override boto.client to return the mock client
boto3.client = lambda service, region_name=None: client
return stubber
def mock_cognito_admin_create_user(user, arguments):
_keep_it_real()
client = boto3.real_client("cognito-idp")
params_admin_create_user = {
"UserPoolId": MOCK_COGNITO_USER_POOL_ID,
"Username": user["email"],
**arguments,
}
stubber = Stubber(client)
# Add responses
stub_response_cognito_list_user_pools(stubber)
stub_response_cognito_admin_create_user(stubber, user, params_admin_create_user)
stubber.activate()
# override boto.client to return the mock client
boto3.client = lambda service, region_name=None: client
return stubber
def mock_sqs(queue_url, event, mock_sqs_send_message_response):
""" Mock SQS send message """
_keep_it_real()
region = "eu-west-2"
client = boto3.real_client("sqs", region_name=region)
stubber = Stubber(client)
event_json = event.to_json()
function_params = {"QueueUrl": queue_url, "MessageBody": event_json}
stubber.add_response(
"send_message", mock_sqs_send_message_response, function_params
)
stubber.activate()
# override boto.client to return the mock client
boto3.client = lambda service, region_name: client
return stubber
def mock_s3_list_objects(bucket_name, prefixes):
_keep_it_real()
client = boto3.real_client("s3")
stubber = Stubber(client)
for prefix in prefixes:
stub_response_s3_list_objects_page_1(stubber, bucket_name, prefix)
stub_response_s3_list_objects_page_2(stubber, bucket_name, prefix)
# replace the get_presigned_url so it runs without AWS creds
client.generate_presigned_url = lambda op, Params, ExpiresIn, HttpMethod: fake_url(
Params["Bucket"], Params["Key"]
)
stubber.activate()
# override boto.client to return the mock client
boto3.client = lambda service, region_name, config=None: client
return stubber
def mock_sts():
_keep_it_real()
client = boto3.real_client("sts")
stubber = Stubber(client)
# mock get_function response
mock_caller_identity = {
"UserId":
"ABCDEFGHIJKLMNOPQRST:1234567890123456789",
"Account":
"123456789012",
"Arn":
"arn:aws:sts::123456789012:assumed-role/IamRoleName/1234567890123456789",
}
stubber.add_response("get_caller_identity", mock_caller_identity)
stubber.activate()
# override boto.client to return the mock client
boto3.client = lambda service: client
return stubber
def mock_sns_publish(sns_event, topic_arn, subject, response_id):
_keep_it_real()
client = boto3.real_client("sns")
stubber = Stubber(client)
# mock get_function response
mock_sns_publish_response = {"MessageId": response_id}
stubber.add_response(
"publish",
mock_sns_publish_response,
{
"TopicArn": topic_arn,
"Message": sns_event,
"Subject": subject,
"MessageStructure": "json",
},
)
stubber.activate()
# override boto.client to return the mock client
boto3.client = lambda service: client
return stubber