🐍

[Python×Lambda] CloudFrontのキャッシュ削除のソースコードとテストコード

2024/06/19に公開

Python

概要

Lambda（Python）で、CloudFrontのキャッシュ削除するための、ソースコードとそのテストコード（覚え書き）

関数のパフォーマンスを向上させるため、Lambdaのベストプラクティスに沿って、実行環境の再利用を活用して関数ハンドラー外で SDK クライアントを宣言しています

コード

言語: Python3
テストライブラリ: Pytest, moto

ソースコード

lambda_function.py

import boto3
import os
from datetime import datetime

# CloudFrontクライアントとディストリビューションIDの初期化
cloudfront = boto3.client('cloudfront')
distribution_id = os.environ.get('DISTRIBUTION_ID')

def lambda_handler(event, context):
    # 無効化バッチの設定
    invalidation_batch = {
        'CallerReference': f'invalidate-{context.aws_request_id}',
        'Paths': {
            'Quantity': 1,   # 無効化するパスの数
            'Items': ['/*']  # 無効化するパスのリスト
        }
    }

    # 無効化リクエストの送信
    response = cloudfront.create_invalidation(
        DistributionId=distribution_id,
        InvalidationBatch=invalidation_batch
    )

    # `Invalidation`オブジェクトから必要な情報を抽出
    invalidation_id = response['Invalidation'].get('Id')
    invalidation_status = response['Invalidation'].get('Status')
    create_time = response['Invalidation'].get('CreateTime')
    create_time_iso = create_time.isoformat() if create_time else None

    # レスポンスを返却
    return {
        'StatusCode': response['ResponseMetadata']['HTTPStatusCode'],
        'InvalidationId': invalidation_id,
        'InvalidationStatus': invalidation_status,
        'CreateTime': create_time_iso
    }

テストコード

test_lambda_function.py

import os
import pytest
from moto import mock_aws
import boto3
from datetime import datetime

@pytest.fixture(scope="function")
def aws_credentials():
    """Mocked AWS credentials for moto."""
    os.environ["AWS_ACCESS_KEY_ID"] = "testing"
    os.environ["AWS_SECRET_ACCESS_KEY"] = "testing"
    os.environ["AWS_SECURITY_TOKEN"] = "testing"
    os.environ["AWS_SESSION_TOKEN"] = "testing"
    os.environ["AWS_DEFAULT_REGION"] = "us-east-1"

@pytest.fixture(scope="function")
def s3_bucket():
    """Mocked S3 bucket using moto."""
    with mock_aws():
        s3 = boto3.client('s3', region_name='us-east-1')
        bucket_name = 'test-bucket'
        s3.create_bucket(Bucket=bucket_name)
        yield bucket_name

@pytest.fixture(scope="function")
def cloudfront_client(aws_credentials, s3_bucket):
    """Mocked CloudFront client using moto."""
    with mock_aws():
        client = boto3.client('cloudfront', region_name='us-east-1')

        # CloudFrontディストリビューションの作成（モック化したS3バケットを使用）
        distribution_config = {
            'CallerReference': 'unique-caller-reference',
            'Origins': {
                'Quantity': 1,
                'Items': [
                    {
                        'Id': 'origin-1',
                        'DomainName': f'{s3_bucket}.s3.amazonaws.com',

                        'OriginPath': '',
                        'CustomHeaders': {'Quantity': 0},
                        'S3OriginConfig': {'OriginAccessIdentity': ''}
                    }
                ]
            },
            'DefaultCacheBehavior': {
                'TargetOriginId': 'origin-1',
                'ViewerProtocolPolicy': 'allow-all',
                'AllowedMethods': {
                    'Quantity': 2,
                    'Items': ['GET', 'HEAD'],
                    'CachedMethods': {'Quantity': 2, 'Items': ['GET', 'HEAD']}
                },
                'ForwardedValues': {
                    'QueryString': False,
                    'Cookies': {'Forward': 'none'},
                    'Headers': {'Quantity': 0},
                    'QueryStringCacheKeys': {'Quantity': 0}
                },
                'MinTTL': 0,
                'DefaultTTL': 86400,
                'MaxTTL': 31536000,
                'Compress': False,
                'LambdaFunctionAssociations': {'Quantity': 0}
            },
            'Comment': '',
            'Enabled': True
        }

        # モック化ディストリビューションの作成
        response =client.create_distribution(DistributionConfig=distribution_config)

        # 取得したディストリビューションIDを環境変数に設定
        os.environ['DISTRIBUTION_ID'] = response['Distribution']['Id']

        yield client

def test_lambda_handler(cloudfront_client):
    """Test the lambda_handler function."""
    # Lambdaハンドラーのimport
    from lambda_function import lambda_handler
    # Lambdaハンドラー用のイベントとコンテキスト
    event = {}
    context = type('Context', (object,), {'aws_request_id': '12345'})()

    # create_invalidation レスポンスのモック化
    mock_response = {
        'ResponseMetadata': {
            'HTTPStatusCode': 201
        },
        'Invalidation': {
            'Id': 'I0123456789',
            'Status': 'InProgress',
            'CreateTime': datetime.now()
        }
    }

    # モックされたcreate_invalidationメソッド
    cloudfront_client.create_invalidation = lambda DistributionId,InvalidationBatch: mock_response

    # 実際のcloudfrontクライアントをモックされたクライアントに置き換え
    original_cloudfront_client = lambda_handler.__globals__['cloudfront']
    lambda_handler.__globals__['cloudfront'] = cloudfront_client

    try:
        # lambda_handlerの呼び出し
        response = lambda_handler(event, context)

        # レスポンスの検証
        assert response['StatusCode'] == 201
        assert response['InvalidationId'] == 'I0123456789'
        assert response['InvalidationStatus'] == 'InProgress'
        assert 'CreateTime' in response
        assert isinstance(response['CreateTime'], str) 
    finally:
        # テスト後に元のclientに戻す
        lambda_handler.__globals__['cloudfront'] = original_cloudfront_client

ポイント

import（from lambda_function import lambda_handler）をテスト関数の中で実行すること
実際（lambda_function.py）のcloudfrontクライアントをモックされたクライアントに置き換えること

参考

※ChatGPT（4o）のお力をかなりお借りしました。

概要

コード

ソースコード

テストコード

ポイント

参考

Discussion