1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
|
import io
import json
import zipfile
from unittest import SkipTest
from unittest.mock import patch
import boto3
from botocore.exceptions import ClientError
from moto import mock_aws, settings
secret_steps = []
def mock_lambda_invoke(*args, **kwarg):
secret_steps.append(json.loads(kwarg["body"])["Step"])
return "n/a"
@mock_aws(config={"lambda": {"use_docker": False}})
@patch(
"moto.awslambda_simple.models.LambdaSimpleBackend.invoke", new=mock_lambda_invoke
)
def test_simple_lambda_is_invoked():
if not settings.TEST_DECORATOR_MODE:
raise SkipTest("Can only test patched code in DecoratorMode")
sm_client = boto3.client("secretsmanager", region_name="us-east-1")
secret_arn = sm_client.create_secret(Name="some", SecretString="secret")["ARN"]
lambda_res = create_mock_rotator_lambda()
sm_client.rotate_secret(
SecretId=secret_arn,
RotationLambdaARN=lambda_res["FunctionArn"],
RotationRules={"AutomaticallyAfterDays": 1, "Duration": "1h"},
RotateImmediately=True,
)
assert secret_steps == ["createSecret", "setSecret", "testSecret", "finishSecret"]
secret_steps.clear()
@mock_aws(config={"lambda": {"use_docker": False}})
@patch(
"moto.awslambda_simple.models.LambdaSimpleBackend.invoke", new=mock_lambda_invoke
)
def test_simple_lambda_is_invoked__do_not_rotate_immediately():
if not settings.TEST_DECORATOR_MODE:
raise SkipTest("Can only test patched code in DecoratorMode")
sm_client = boto3.client("secretsmanager", region_name="us-east-1")
secret_arn = sm_client.create_secret(Name="some", SecretString="secret")["ARN"]
lambda_res = create_mock_rotator_lambda()
sm_client.rotate_secret(
SecretId=secret_arn,
RotationLambdaARN=lambda_res["FunctionArn"],
RotationRules={"AutomaticallyAfterDays": 1, "Duration": "1h"},
RotateImmediately=False,
)
assert secret_steps == ["testSecret"]
secret_steps.clear()
def mock_lambda_zip():
code = """
def lambda_handler(event, context):
return event
"""
zip_output = io.BytesIO()
zip_file = zipfile.ZipFile(zip_output, "w", zipfile.ZIP_DEFLATED)
zip_file.writestr("lambda_function.py", code)
zip_file.close()
zip_output.seek(0)
return zip_output.read()
def create_mock_rotator_lambda():
client = boto3.client("lambda", region_name="us-east-1")
return client.create_function(
FunctionName="mock-rotator",
Runtime="python3.9",
Role=get_mock_role_arn(),
Handler="lambda_function.lambda_handler",
Code={"ZipFile": mock_lambda_zip()},
)
def get_mock_role_arn():
iam = boto3.client("iam", region_name="us-east-1")
try:
return iam.get_role(RoleName="my-role")["Role"]["Arn"]
except ClientError:
return iam.create_role(
RoleName="my-role",
AssumeRolePolicyDocument="some policy",
Path="/my-path/",
)["Role"]["Arn"]
|
