test_scheduled_lambda.py - AWS Code Sample

test_scheduled_lambda.py

# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. # SPDX-License-Identifier: Apache-2.0 """ Unit tests for scheduled_lambda.py functions. """ import boto3 from botocore.exceptions import ClientError import pytest import scheduled_lambda @pytest.mark.parametrize('failed_target_count,error_code,stop_on_method', [ (0, None, None), (1, None, None), (0, 'TestException', 'stub_put_rule'), (0, 'TestException', 'stub_add_permission'), (0, 'TestException', 'stub_put_targets') ]) def test_schedule_lambda_function( make_stubber, stub_runner, failed_target_count, error_code, stop_on_method): eventbridge_client = boto3.client('events') eventbridge_stubber = make_stubber(eventbridge_client) lambda_client = boto3.client('lambda') lambda_stubber = make_stubber(lambda_client) event_rule_name = 'test-rule' event_schedule = 'test-schedule' event_rule_arn = f'arn:aws:events:::rules/{event_rule_name}' lambda_func_name = 'test-func' lambda_func_arn = f'arn:aws:lambda:::functions/{lambda_func_name}' with stub_runner(error_code, stop_on_method) as runner: runner.add( eventbridge_stubber.stub_put_rule, event_rule_name, event_schedule, event_rule_arn) runner.add( lambda_stubber.stub_add_permission, lambda_func_name, 'lambda:InvokeFunction', 'events.amazonaws.com', event_rule_arn) runner.add( eventbridge_stubber.stub_put_targets, event_rule_name, [{'Id': lambda_func_name, 'Arn': lambda_func_arn}], failed_count=failed_target_count) if error_code is None: got_arn = scheduled_lambda.schedule_lambda_function( eventbridge_client, event_rule_name, event_schedule, lambda_client, lambda_func_name, lambda_func_arn) assert got_arn == event_rule_arn else: with pytest.raises(ClientError) as exc_info: scheduled_lambda.schedule_lambda_function( eventbridge_client, event_rule_name, event_schedule, lambda_client, lambda_func_name, lambda_func_arn) assert exc_info.value.response['Error']['Code'] == error_code @pytest.mark.parametrize('enable,error_code', [ (True, None), (False, None), (True, 'TestException') ]) def test_update_event(make_stubber, enable, error_code): eventbridge_client = boto3.client('events') eventbridge_stubber = make_stubber(eventbridge_client) event_rule_name = 'test-rule' if enable: eventbridge_stubber.stub_enable_rule(event_rule_name, error_code=error_code) else: eventbridge_stubber.stub_disable_rule(event_rule_name, error_code=error_code) if error_code is None: scheduled_lambda.update_event_rule(eventbridge_client, event_rule_name, enable) else: with pytest.raises(ClientError) as exc_info: scheduled_lambda.update_event_rule( eventbridge_client, event_rule_name, enable) assert exc_info.value.response['Error']['Code'] == error_code @pytest.mark.parametrize('state,error_code', [ ('ENABLED', None), ('DISABLED', None), ('ENABLED', 'TestException')]) def test_get_event_enabled(make_stubber, state, error_code): eventbridge_client = boto3.client('events') eventbridge_stubber = make_stubber(eventbridge_client) event_rule_name = 'test-rule' eventbridge_stubber.stub_describe_rule( event_rule_name, state, error_code=error_code) if error_code is None: got_enabled = scheduled_lambda.get_event_rule_enabled( eventbridge_client, event_rule_name) assert got_enabled == (state == 'ENABLED') else: with pytest.raises(ClientError) as exc_info: scheduled_lambda.get_event_rule_enabled(eventbridge_client, event_rule_name) assert exc_info.value.response['Error']['Code'] == error_code @pytest.mark.parametrize('error_code,stop_on_method', [ (None, None), ('TestException', 'stub_remove_targets')]) def test_delete_event(make_stubber, stub_runner, error_code, stop_on_method): eventbridge_client = boto3.client('events') eventbridge_stubber = make_stubber(eventbridge_client) event_rule_name = 'test-rule' lambda_func_name = 'test-func' with stub_runner(error_code, stop_on_method) as runner: runner.add( eventbridge_stubber.stub_remove_targets, event_rule_name, [lambda_func_name]) runner.add( eventbridge_stubber.stub_delete_rule, event_rule_name) if error_code is None: scheduled_lambda.delete_event_rule( eventbridge_client, event_rule_name, lambda_func_name) else: with pytest.raises(ClientError) as exc_info: scheduled_lambda.delete_event_rule( eventbridge_client, event_rule_name, lambda_func_name) assert exc_info.value.response['Error']['Code'] == error_code