Zenn
🐍

[Python×Lambda] SNS通知のためのソースコードとテストコード（※後者に重点）

2024/04/08に公開
Python
Lambda
pytest
moto
tech

はじめに

Lambda（Python）で、SNS通知するための、ソースコードとそのテストコード

Lambdaのベストプラクティスによれば、実行環境の再利用を活用して関数のパフォーマンスを向上させるため、関数ハンドラー外で SDK クライアントを宣言すべき、とのことなのですが、、モック（moto）を用いたテストコード実装時に沼ったので、覚え書きしておきます

コード

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

ソースコード

lambda_function.py
import boto3
import os

sns_client = boto3.client("sns")

topic_arn = os.environ["TOPIC_ARN"]
sns_message_type = os.environ["SNS_MESSAGE_TYPE"]

def lambda_handler(event, context):
    """
    eventをチェックして、エラーがあった場合は SNS に通知する
    """

    # SNS通知用の各種変数を初期化
    subject=""
    message=""

    has_error = get_error(event)

    # エラーが見つかった場合 SNS にイベントを発行する
    if any(has_error):
      function_name = context.function_name
     
      env_name = "本番環境" if "prod" in function_name else "テスト環境"
     
      subject = f"エラー通知メール（{env_name}）"
      message = "メッセージ\n本文"
     
      sns_client.publish(
            TopicArn=topic_arn,
            Subject=subject,
            Message=message,
            MessageAttributes={
                "type": {
                    "DataType": "String",
                    "StringValue": sns_message_type
                }
            }
      )

    # テストコード検証用: 特に実環境には影響なし
    return {"TopicArn":topic_arn, "Subject": subject, "Message": message}
   
# 以下、省略めにしてます
def get_error(event):
    # エラーなしの場合
    if xxxxx
        return {}
    # エラーありの場合
    return {
        "error_message":"xxxx"
    }

テストコード

test_lambda_function.py
import pytest
from moto import mock_aws

import boto3
import os

class MockContext:
    """
    関数名設定用
    """
    def __init__(self, function_name):
        self.function_name = function_name

@pytest.fixture(scope="function")
def aws_credentials():
    """
    AWS認証情報のモック設定
    """
    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 sns_setup(aws_credentials):
    """
    環境変数TOPIC_ARNとSNSクライアントをモック化
    """
    with mock_aws():
        sns_client = boto3.client("sns", region_name="us-east-1")
        response = sns_client.create_topic(Name="test-topic")
        topic_arn =  response["TopicArn"]
        os.environ["TOPIC_ARN"] =  response["TopicArn"]
        yield topic_arn

def test_lambda_handler_no_error(sns_setup):
    """
    エラーなしのテスト
    """
    from lambda_function import lambda_handler
    event = {}
    context = MockContext(function_name = "error-finder-test")
    result = lambda_handler(event, context)
    # 結果確認
    assert result["TopicArn"] == "arn:aws:sns:us-east-1:123456789012:test-topic"
    assert result["Subject"] == ""
    assert result["Message"] == ""

def test_lambda_handle_error(sns_setup):
    """
    エラー通知のテスト
    """
    from lambda_function import lambda_handler
    event = {}
    context = MockContext(function_name = "error-finder-prod")
    result = lambda_handler(event, context)
    # 結果確認
    assert result["TopicArn"] == "arn:aws:sns:us-east-1:123456789012:test-topic"
    assert result["Subject"] == "タイトル"
    assert result["Message"] == "メッセージ\n本文"

ポイント

ポイントは、import（from lambda_function import lambda_handler）をテスト関数の中で実行すること
コチラのブログが大いに参考になりました！

補足

AWS公式によれば、SNSで設定できるメールタイトル（subject）は100文字までらしいです。
後日対応でタイトルを長くしたところ、テスト中にエラーが発生して焦りました、、

参考

Discussion