def test_lambda_handler_with_valid_instance_id():
param = '/ami-baking-lnx-amzn-soe/lnx-amzn/instanceId'
# Setup test data
original_param_value = 'blank' # same as initial cfn deployment value
client = boto3.client('ssm', region_name=CONST_REGION)
client.put_parameter(
Name=param,
Value=original_param_value,
Type='String',
Overwrite=True
)
# Test Lambda handler
next_ami_id = "ami-12345678"
instance_id = 'i-87654321'
event = {"AMI": next_ami_id, "BuildInstanceID": instance_id}
context = ContextMock()
output_event = lambda_handler(copy.deepcopy(event), context)
# Verify the output event
assert output_event != event # Assert that the event has been modified
assert output_event["SsmParamVersion"] == '1'
assert output_event["SsmParam"] == CONST_NEXT_AMI_PARAM
# Verify the SSM param is actually updated
param_value = get_ssm_param(CONST_REGION, param)
assert param_value != original_param_value
assert param_value == instance_id
def test_lambda_handler():
"""Test update_next_ami_function.lambda_handler to update next"""
# Setup test data
original_param_value = 'blank' # same as initial cfn deployment value
client = boto3.client('ssm', region_name=CONST_REGION)
client.put_parameter(
Name=CONST_NEXT_AMI_PARAM,
Value=original_param_value,
Type='String',
Overwrite=True
)
# Test Lambda handler
next_ami_id = "ami-1234567890abcdef"
event = {"AMI": next_ami_id, "BuildInstanceID": 'i-12345678'}
context = ContextMock()
output_event = lambda_handler(copy.deepcopy(event), context)
# Verify the output_event is updated with the param store path and its version
assert output_event != event # Assert that the event has been modified
assert output_event["SsmParamVersion"] == '2'
assert output_event["SsmParam"] == CONST_NEXT_AMI_PARAM
# Verify the SSM param is actually updated
param_value = get_ssm_param(CONST_REGION, CONST_NEXT_AMI_PARAM)
assert param_value != original_param_value
assert param_value == next_ami_id
def test_lambda_handler_new_roles(test_role_creation, account_ids,
expected_error):
"""Test get_accounts_function.lambda_handler to return all account ids"""
iam_client = boto3.client('iam', region_name=CONST_REGION)
role_name = (CONST_SOL_NAMING + "-" + CONST_SOE_TYPE +
"-soe-service-iam-role")
if test_role_creation:
# Verify role does not exist
with pytest.raises(ClientError) as excinfo:
iam_client.get_role(RoleName=role_name)
assert 'Role {} not found'.format(role_name) in str(excinfo.value)
else:
# Create the role before the lambda_handler to test the update
init_policy_doc = json.dumps({
"Version": "2012-10-17",
"Statement": []
})
role = iam_client.create_role(Path='/service/',
RoleName=role_name,
AssumeRolePolicyDocument=init_policy_doc,
Description='Unit Test initialised role')
print('role')
print(role)
role_response = iam_client.get_role(RoleName=role_name)
assert role_response['Role']['RoleName'] == role_name
role_policy_statements = role_response['Role'][
'AssumeRolePolicyDocument']['Statement']
assert not role_policy_statements # set to [] in above policy doc
# Inputs for Lambda handler
event = {'Accounts': account_ids}
context = ContextMock()
# Test Lambda handler
if expected_error:
with pytest.raises(expected_error) as excinfo:
lambda_handler(event, context)
assert "'NoneType' object is not iterable" in str(excinfo.value)
else:
output_event = lambda_handler(copy.deepcopy(event), context)
# Assert that the event is not modified
assert output_event == event
# Verify the role has been created/updated correctly
role_response = iam_client.get_role(RoleName=role_name)
assert role_response['Role']['RoleName'] == role_name
role_policy_statements = role_response['Role'][
'AssumeRolePolicyDocument']['Statement']
assert len(role_policy_statements) == 1
assert role_policy_statements[0]['Action'] == 'sts:AssumeRole'
assert role_policy_statements[0]['Effect'] == 'Allow'
principal_accounts = [
"arn:aws:iam::" + account_id + ":root"
for account_id in account_ids
]
assert role_policy_statements[0]['Principal'][
'AWS'] == principal_accounts
def test_lambda_handler_with_invalid_instance_id():
# Test Lambda handler
next_ami_id = "ami-12345678"
event = {"AMI": next_ami_id, "BuildInstanceID": None}
context = ContextMock()
with pytest.raises(ValueError) as excinfo:
lambda_handler(event, context)
assert 'BuildInstanceID is None' in str(excinfo.value)
def test_lambda_handler_with_missing_instance_id():
"""Test update_next_ami_function.lambda_handler to update next"""
# Test Lambda handler
next_ami_id = "ami-12345678"
event = {"AMI": next_ami_id}
context = ContextMock()
with pytest.raises(KeyError) as excinfo:
lambda_handler(event, context)
assert 'BuildInstanceID' in str(excinfo.value)
def test_lambda_handler_with_invalid_ami_id():
"""Test update_next_ami_function.lambda_handler to update next"""
# Test Lambda handler
next_ami_id = "ami-invalid-id"
event = {"AMI": next_ami_id, "BuildInstanceID": 'i-12345678'}
context = ContextMock()
with pytest.raises(ValueError) as excinfo:
lambda_handler(event, context)
assert 'does not match AMI Id format' in str(excinfo.value)
def test_lambda_handler(mock_boto_client, monkeypatch):
"""Test trigger_test_function.lambda_handler"""
# Setup mock input
mock_build_automation_execution_id = 'mock_build_automation_execution_id'
mock_instance_id = 'i-12341234'
ssm_document = 'SSMDocument'
# Set Environment Variables
monkeypatch.setenv("SSMDocument", ssm_document)
# We need to reload the module so it re-reads the os.environ values we set above
reload(trigger_test_function)
# patch trigger_ssm as moto does NOT support start_automation_execution yet
# https://github.com/spulec/moto/blame/603f7c58a230919da3ee836575351366e46cc26c/IMPLEMENTATION_COVERAGE.md#L4022
# mock_boto_client = MagicMock(return_value=MOCK_TEST_AUTOMATION_EXECUTION_ID)
mock_boto_client.side_effect = mock_ssm_client
# Setup Input
event = {
'BuildAutomationExecutionId': mock_build_automation_execution_id,
'InstanceID': mock_instance_id
}
context = ContextMock()
# Test Lambda handler
output_event = trigger_test_function.lambda_handler(
copy.deepcopy(event), context)
print("output_event")
print(output_event)
# Verify the output
assert output_event != event # Assert that the event has been modified
assert output_event["BuildInstanceID"] == mock_instance_id
assert output_event[
"TestAutomationExecutionId"] == MOCK_TEST_AUTOMATION_EXECUTION_ID
assert output_event["AMI"] == MOCK_AMI_ID
# Assert the boto client mock is called exactly as expected
assert mock_boto_client.call_count == 2
expected = [
call('GetAutomationExecution',
{'AutomationExecutionId': mock_build_automation_execution_id}),
call(
'StartAutomationExecution', {
'DocumentName': ssm_document,
'Parameters': {
'sourceAMIid': [MOCK_AMI_ID]
}
})
]
assert mock_boto_client.call_args_list == expected
def test_lambda_handler(mock_get_automation_execution,
automation_execution_status, build_status):
"""Test check_build_function.lambda_handler"""
# Setup mock input
mock_state = automation_execution_status
mock_instance_id = 'i-12345678'
mock_automation_execution_id = 'mock_automation_execution_id'
# patch get_automation_execution as moto does NOT support it yet
mock_get_automation_execution.side_effect = MagicMock(
return_value={
'AutomationExecution': {
'AutomationExecutionStatus': mock_state,
'Outputs': {
'startInstances.InstanceIds':
[" \"%s\" " % mock_instance_id]
}
}
})
# Setup Input
event = {'BuildAutomationExecutionId': mock_automation_execution_id}
context = ContextMock()
# Test Lambda handler
output_event = lambda_handler(copy.deepcopy(event), context)
# Verify the output event
assert output_event != event # Assert that the event has been modified
assert output_event["InstanceID"] == mock_instance_id
assert output_event["BuildStatus"] == build_status
assert output_event["CheckType"] == "ssm_build"
# Verify the mocks have been called as expected
# TODO: Update lambda_handler to call once instead of twice the exact same request
assert mock_get_automation_execution.call_count == 2
expected = [
call('GetAutomationExecution',
{'AutomationExecutionId': mock_automation_execution_id}),
call('GetAutomationExecution',
{'AutomationExecutionId': mock_automation_execution_id}),
]
assert mock_get_automation_execution.call_args_list == expected
def test_lambda_handler(keep_ec2, expect_termination):
"""Test get_accounts_function.lambda_handler to return all account ids"""
# Setup a test instance to delete
ec2_client = boto3.client('ec2', region_name=CONST_REGION)
run_ec2_response = ec2_client.run_instances(ImageId='ami-12345678',
MaxCount=1,
MinCount=1)
instance_id = run_ec2_response['Instances'][0]['InstanceId']
print(instance_id)
# Check the instance is initially running
response = ec2_client.describe_instances(InstanceIds=[instance_id])
instance_state = response['Reservations'][0]['Instances'][0]['State'][
'Name']
assert instance_state == 'running'
# Test Lambda handler
event = {'InstanceID': instance_id}
context = ContextMock()
if keep_ec2:
event['KeepTestInstance'] = keep_ec2
output_event = lambda_handler(copy.deepcopy(event), context)
# Assert that the event is not modified
assert output_event == event
# Check the instance is terminated (moto is instant vs boto will take time to terminate)
response = ec2_client.describe_instances(InstanceIds=[instance_id])
instance_state = response['Reservations'][0]['Instances'][0]['State'][
'Name']
if expect_termination:
assert instance_state == 'terminated'
else:
assert instance_state == 'running'
def test_lambda_handler(ami_pattern, pipeline_override_ami, event_override_ami,
expected_error, exception_message, monkeypatch):
"""Test trigger_build_function.lambda_handler"""
# Setup mock input
soe_image_name = 'plt-baking-soe-ami-for-unit-test'
mock_automation_execution_id = 'mock_automation_execution_id'
ssm_document = 'SSMDocument'
soe_type = 'lnx'
ami_owner = 'self'
# Set Environment Variables
monkeypatch.setenv("SSMDocument", ssm_document)
monkeypatch.setenv("SOEType", soe_type)
monkeypatch.setenv("AMIPattern", ami_pattern)
monkeypatch.setenv("AMIOwner", ami_owner)
monkeypatch.setenv("OverrideAMI", pipeline_override_ami)
# We need to reload the module so it re-reads the os.environ values we set above
reload(trigger_build_function)
# patch trigger_ssm as moto does NOT support start_automation_execution yet
# https://github.com/spulec/moto/blame/603f7c58a230919da3ee836575351366e46cc26c/IMPLEMENTATION_COVERAGE.md#L4022
mock_trigger_ssm = MagicMock(return_value=mock_automation_execution_id)
monkeypatch.setattr('trigger_build.trigger_build_function.trigger_ssm',
mock_trigger_ssm)
if pipeline_override_ami:
ami_id = pipeline_override_ami
else:
# Setup mock Public AMI (we need to create a AMI and modify launch permissions)
# 1. Create a mock SOE image from a instance
ec2_client = boto3.client('ec2', region_name=CONST_REGION)
run_ec2_response = ec2_client.run_instances(ImageId='ami-12345678',
MaxCount=1,
MinCount=1)
instance_id = run_ec2_response['Instances'][0]['InstanceId']
print(instance_id)
create_ami_response = ec2_client.create_image(
Description='Moto unit test image 1',
InstanceId=instance_id,
Name=soe_image_name)
ami_id = create_ami_response['ImageId']
# 2. Make mock SOE image Public
exec_user_id = '123456789012'
# We need to override the EXEC_USERS in the actual module as moto does not support 'all' yet
# https://github.com/spulec/moto/blob/d8dbc6a49ccf969f50ed3f7b52f341db3a7715f0/moto/ec2/models.py#L1190
monkeypatch.setattr('trigger_build.trigger_build_function.EXEC_USERS',
exec_user_id)
modify_image_response = ec2_client.modify_image_attribute(
Attribute='launchPermission',
ImageId=ami_id,
OperationType='add',
UserIds=[exec_user_id],
UserGroups=['all'])
print("modify_image_response")
print(modify_image_response)
# Setup Input
event = {}
context = ContextMock()
if event_override_ami:
event['OverrideAMI'] = event_override_ami
ami_id = event_override_ami
# Test Lambda handler
if isinstance(expected_error, type):
with pytest.raises(expected_error) as excinfo:
trigger_build_function.lambda_handler(event, context)
#TODO: Improve the error handling in the lambda_handler
assert exception_message in str(excinfo.value)
mock_trigger_ssm.assert_not_called()
else:
output_event = trigger_build_function.lambda_handler(
copy.deepcopy(event), context)
print("output_event")
print(output_event)
# Verify the output
assert output_event != event # Assert that the event has been modified
assert output_event[
"BuildAutomationExecutionId"] == mock_automation_execution_id
mock_trigger_ssm.assert_called_with(CONST_SOL_NAMING, CONST_REGION,
ssm_document, soe_type, ami_id)