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)