AWS Lambdaでpytestする際のメモ(fixture/mock)

AWS Lambda(+API Gateway)を利用したアプリケーションを作成する際に、よく使うPytestのFixtureやモックの使い方を自分向けにメモしておきます。

Pytestの書き方はいつも使っていないとすぐに忘れてしまいます。

本記事の内容

1. 前提
2. Lambdaプログラム
3. 普通にテスト
4. jsonファイルからAPI Gateway Requestを読み込んで利用する
5. lambda_handler内で呼び出している関数をモックする
6. lambda_handler内で呼び出している関数をモックする(json使わない例)
- 6.1. 複数パラメータでのテスト
Pythonのオススメ勉強方法

1. 前提

PythonやPytestのセットアップなどは完了済みとします。
今回はAmazonの商品コードであるasinを渡すとその商品情報を取得して来るというアプリを想定してテストしています。
プログラムファイルとしてはlambda_frontディレクトリ内のmain.pyとなります。

2. Lambdaプログラム

プログラムについてです。いくつかの処理は省略してますが、大体以下のような感じです。

invoke_lambda関数では非同期処理にしたかったのでスクレイピング用のLambda関数を呼び出してます。

成功したらreturn_successを、失敗したらreturn_errorを返します。

def lambda_handler(event: dict, context):

    if event['httpMethod'] == 'GET' and event['path'] == "/asin":
        invoke_lambda(event)
        return return_success('Received your requests. please wait...')

    return return_error("Invalid Error!!!")

def return_success(data):
    return {
        'statusCode': 200,
        'body': json.dumps(data)
    }

def return_error(data):
    return {
        'statusCode': 501,
        'body': json.dumps(data)
    }

def invoke_lambda(event):
    <<スクレイピングを実行するLambdaを呼び出すコード>>

3. 普通にテスト

まずは普通にテストしてみます。簡単にreturn_successのテストをしてみます。

import json
import pytest
import lambda_front.main as front

def test_return_success():
    res = front.return_success('test')
    assert res["body"] == json.dumps('test')
    assert res["statusCode"] == 200

4. jsonファイルからAPI Gateway Requestを読み込んで利用する

conftest.pyを作成してから、jsonファイルを読み込んでPytestのfixtureとして利用します。

以下がconftest.pyの内容です。単純にjsonファイルを読み込んでるだけですね。

import pytest
import json

@pytest.fixture()
def api_pxy_data():
    with open('data/api_pxy.json') as f:
        test_api_event = json.load(f)
        print("before test_api_event=", test_api_event)
        yield test_api_event
        print("after test_api_event")

テストしたいAPI GatewayからのRequestをjsonで用意しておきます。今回のjsonファイルでは以下のリクエストを想定しています。

https://<api url>/prod/asin/asin=B08NYC3LFJ

{
    "resource": "/{proxy+}",
    "path": "/asin",
    "httpMethod": "GET",
    "headers": {

(省略)

    "queryStringParameters": {
        "asin": "B08NYC3LFJ"
    },

    "multiValueQueryStringParameters": {
        "asin": [
            "B08NYC3LFJ"
        ]
    },

(省略)

    },
    "body": null,
    "isBase64Encoded": false

実際にtest_main.pyからapi_pxy_dataを呼び出してみます。

def test_lambda_handler(api_pxy_data):
    res = front.lambda_handler(event=api_pxy_data, context={})
    assert res['statusCode'] == 200

5. lambda_handler内で呼び出している関数をモックする

実は先ほどのテストは失敗します。 lambda_handler内で読んでいるinvoke_lambda(event)では別のLambdaを呼び出しているのですが、ローカルPCからは実行できないためです。

そこでinvoke_lambdaが成功するようにモックを使います。 mocker.patchを使ってinvoke_lambda関数を呼び出した時に強制的にTrueを返すようにして実際に実行しないようにします。

from unittest import mock

def test_lambda_handler(api_pxy_data, mocker):
    mocker.patch("lambda_front.main.invoke_lambda", return_value=True)

    res = front.lambda_handler(event=api_pxy_data, context={})
    assert front.invoke_lambda(0) == True
    assert res['statusCode'] == 200

こんな感じで書く事もできます。こっちの方が見やすいですね。

@mock.patch("lambda_front.main.invoke_lambda", return_value=True)
def test_lambda_handler(mock_invoke, api_pxy_data):
    # mocker.patch("lambda_front.main.invoke_lambda", return_value=True)

    res = front.lambda_handler(event=api_pxy_data, context={})
    assert front.invoke_lambda(0) == True
    mock_invoke.assert_called()
    assert res['statusCode'] == 200

6. lambda_handler内で呼び出している関数をモックする(json使わない例)

jsonを使う場合には複数ケースで実行するテストがしにくいです。なので、一旦以下のような形で必要なパラメータだけを渡してテストする形にします。

@mock.patch("lambda_front.main.invoke_lambda", return_value=True)
def test_lambda_handler_event(mock_invoke):
    event = {}
    event["path"] = "/asin"
    event["httpMethod"] = "GET"
    assert front.invoke_lambda(0) == True
    mock_invoke.assert_called()
    res = front.lambda_handler(event, context={})
    assert res['statusCode'] == 200

6.1. 複数パラメータでのテスト

最終的に複数パラメータを指定してテストを行います。以下の例では/asinでGETの場合には200を返していること、それ以外は501を返すかをテストしています。

@pytest.mark.parametrize('event_data, status_code',
                         [
                             ({"path": '/asin', "httpMethod": 'GET'}, 200),
                             ({"path": '/test', "httpMethod": 'GET'}, 501),
                             ({"path": '/asin', "httpMethod": 'POST'}, 501)
                         ])

@mock.patch("lambda_front.main.invoke_lambda", return_value=True)
def test_lambda_handler_multi_event(mock_invoke, event_data, status_code):
    res = front.lambda_handler(event_data, context={})
    assert res['statusCode'] == status_code