【pytest】AWS IoT メッセージを Publish する Lambda 関数をローカルでテストしたい

# 1. 前書き
AWS IoT メッセージを Publish をする Lambda 関数と関数を呼び出すための API Gateway を CDK で作成しました。
構成は以下のようになっています。

```
sample/
 ├ bin/
    └ sample.ts
 ├ lib/
    └ sample-stack.ts
 ├ iot/
    └ publish.py
 ├
...以下略
```

CDK で使用した言語は TypeScript 、Lambda 関数で使用している言語は python です。

```python: publish.py
# coding: utf-8
import os
import json
import logging
import boto3

# 環境変数から取得
REGION = os.environ.get('REGION')

# AWS IoT Dataオブジェクトを取得
iot = boto3.client('iot-data', region_name=REGION)

# トピック名
TOPIC_NAME = 'sample/pub'

# Lambdaのメイン関数
def handler(event, context):

    # レスポンスの雛形
    response = {
        'statusCode': 200,
        'body': '',
        'headers': {
            'Content-Type': 'application/json',
            'Access-Control-Allow-Origin': '*',
        },
    }
    # メッセージの内容
    payload = {
        "test": {
            "name":"sample"
        }
    }

    try:
        # メッセージをPublish
        iot.publish(
            topic=TOPIC_NAME,
            qos=0,
            payload=json.dumps(payload, ensure_ascii=False)
        )

        logging.info('Succeeded.')
        response['body'] = json.dumps("Payload: {0}".format(payload))
        response['statusCode'] = 200

    except Exception as e:
        logging.info('Failed.')
        response['body'] = json.dumps("Error Message: {0}".format(e))
        response['statusCode'] = 500

    return response
```

<details><summary>他のソースコード</summary>

```typescript: sample.ts
#!/usr/bin/env node
import * as cdk from 'aws-cdk-lib';
import { SampleStack } from '../lib/sample-stack';

const app = new cdk.App();
new SampleStack(app, 'SampleStack');
```

```typescript: sample-stack.ts
import * as cdk from 'aws-cdk-lib';
import { Construct } from 'constructs';
import * as iam from 'aws-cdk-lib/aws-iam'
import * as lambda from 'aws-cdk-lib/aws-lambda';
import * as apigateway from 'aws-cdk-lib/aws-apigateway'

export class SampleStack extends cdk.Stack {
  constructor(scope: Construct, id: string, props?: cdk.StackProps) {
    super(scope, id, props);

    const ACCOUNT = cdk.Stack.of(this).account;
    const REGION = cdk.Stack.of(this).region;
    const TOPIC_NAME = 'sample'

    // IAM Role
    // publish用のロール
    const postRole = new iam.Role(this, 'postRole',{
        assumedBy: new iam.ServicePrincipal('lambda.amazonaws.com'),
        description: 'post iot publish',
        roleName: 'post-iot-publish-role'
    });

    /*
    * リソースの生成
    */

    // Lambda
    const postFunction = new lambda.Function(this, 'postFunction',{
        functionName: 'postIotPublish',
        runtime: lambda.Runtime.PYTHON_3_10,
        code: lambda.Code.fromAsset('iot'),
        handler: 'publish.handler',
        environment: {
            REGION: REGION,
        },
        role: postRole,
    });

    // API Gateway
    const api = new apigateway.RestApi(this, 'ApiGateway', {
        restApiName: 'apigateway',
        description: 'Publish用のAPIGateway',
        endpointTypes: [apigateway.EndpointType.REGIONAL],
    });

    const test = api.root.addResource('test', {
        defaultCorsPreflightOptions: {
            allowOrigins: apigateway.Cors.ALL_ORIGINS,
            allowMethods: apigateway.Cors.ALL_METHODS,
            allowHeaders: apigateway.Cors.DEFAULT_HEADERS,
            statusCode: 200,
        }
    })

    const postMethod = test.addResource('{id}').addMethod('POST',new apigateway.LambdaIntegration(postFunction),{
        methodResponses:[{
            statusCode: '200',
            responseModels: {
                'application/json': apigateway.Model.EMPTY_MODEL,
            }
        }]
    });
    // stage
    const deployment = new apigateway.Deployment(this, 'Deployment', {
        api:api,
    });

    // 使用量プラン
    // 使用量プランを作成してステージに紐付け
    const plan = api.addUsagePlan('usagePlan', {
        name: 'usagePlan',
        throttle: {
            burstLimit: 5000, // 1~数秒間の最大APIリクエストレート制限
            rateLimit: 10000, // API リクエストの定常状態レート制限 (長期間にわたる1秒あたりの平均リクエスト)
        },
        quota: {
            limit: 10000, // 指定された期間内にユーザが実行できるリクエストの最大数
            period: apigateway.Period.MONTH, // リクエストの最大制限が適用される期間
        }
    });

    plan.addApiStage({ stage: api.deploymentStage });

    // ポリシー
    const postPolicyStatement = new iam.PolicyStatement({
        effect: iam.Effect.ALLOW,
        actions: ["iot:Publish"],
        resources: [`arn:aws:iot:${REGION}:${ACCOUNT}:topic/${TOPIC_NAME}/*`]
    })

    const postPolicy = new iam.ManagedPolicy(this, 'postPolicy', {
        managedPolicyName: "postPolicy",
        description: 'SNSトピックにメッセージをPublishする権限',
        statements: [postPolicyStatement]
    })

    const lambdaBasicExecutionManagedPolicy = iam.ManagedPolicy.fromAwsManagedPolicyName(
        'service-role/AWSLambdaBasicExecutionRole',
    )

    postRole.addManagedPolicy(lambdaBasicExecutionManagedPolicy);
    postRole.addManagedPolicy(postPolicy);
  }
}
```

</details>

API Gateway を通じて Lambda が実行され、該当の MQTT トピックに対してメッセージを Publish します。レスポンスとして API Gateway に成功・失敗を返します。

この Lambda 関数がうまく動作するかどうかテストをするために単体テストを行います。 python のテストツールである pytest を用いてテストをしようとしました。しかし、このまま実行しても実際の AWS リソースにアクセスしようとしてしまうため、うまくテストができませんでした。このような場合はモックをすることで AWS 環境を擬似的に用意する必要があるということで実際にやってみました。

# 2. 本題
### 使用したライブラリ
モックをするために使用したライブラリが`pytest-mock`と`moto`です。
`pytest-mock`は既存ライブラリや関数等をモックするプラグインライブラリです。 pytest 時にモックされたオブジェクトの実挙動を回避してテストすることができます。

```
pip install pytest pytest-mock
```

`moto`は AWS の各種サービスをモックするための Python ライブラリです。実際の AWS 環境にアクセスすることなくテストをすることができます。
moto[iotdata] の最新バージョンでは`jsondiff`がインポートされなかったため、`2.0.1`をインストールします。

```
pip install 'moto[iotdata==2.0.1]'
```

### pytestの作成

まずは`__init__.py`を用意します。

```python: __init__.py
from dataclasses import dataclass
import os

@dataclass
class getEnvNames:
    def __init__(self):
        self.REGION = os.environ["REGION"]
```

pytestを実行する上で必要なファイルです。中には環境変数である`REGION`を定義しておきます。

続いてテストファイルを作成します。
この中で`mocker`を用いてAWSリソースをモックしていきます。

```python: test_iot_publish.py
import os
import pytest
import boto3
import logging

from moto import mock_iotdata

MODULE_NAME = 'iot' + '.publish'

@pytest.fixture(scope='module', autouse=True)
def iot_data_client():
    with mock_iotdata():
        iot_data = boto3.client('iot-data', region_name='ap-northeast-1')
    yield iot_data

@pytest.fixture(scope='function', autouse=True)
def scope_session(mocker, iot_data_client):
    mocker.patch.dict(os.environ, {"REGION":"ap-northeast-1"})
    mocker.patch(
        MODULE_NAME + '.iot',
        iot_data_client
    )
    yield

def test_pattern():
    """メッセージをpublishする"""
    logging.debug('トピックにpublishする')
    from .publish import handler

    # API Gatewayからのevent設定
    event_data = {
        'httpMethod': 'POST',
        "resource": "/test/user/{id}"
    }

    with mock_iotdata():
        response = handler(event_data, '')
        assert response['statusCode'] == 200
        assert response['body'] == '"Payload: {\'test\': {\'name\': \'sample\'}}"'
```

`iot_data_client()`では、`mock_iotdata`を用いてiot_dataクライアントをモックしたものを返す関数です。`scope=module`なので、テストファイル単位で実行されます。
`scope_session()`では、`mocker.patch`を使用して、変数やオブジェクトをモックしています。今回は環境変数に入る値とiot_data_clientオブジェクトをモックしています。`scope=function`なので関数ごとに実行されます。
`test_pattern()`はテストの本体です。テストしたい関数をインポートして検証しています。関数を呼び出す際には`mock_iodata`を使用する必要があります。

これでテストファイルの作成ができました。pytestを実行すると問題なく`pass`することが確認できます。

# 3. まとめ

iot、iot-dataをモックする例が見当たらなかったので記事にしました。pytestとmotoを用いて簡単にテストすることができました。
これらを使用するメリットとしては、Lambda関数のテストを行うために関連するリソースを一度デプロイする必要がないという点です。Cognitoなどを使用している場合、認証情報などのやり取りが手間になってしまいます。そのため、Lambda関数のみをテストするのであれば、ローカルで簡単にテストできる`pytest-mock`,`moto`を活用しましょう。