AWS Lambdaの引数validateをdecorator化しつつテストを追加してみた

AWS Lambdaに追加済み関数の引数バリデーションコードが雑多になりがちだったため、関数に分離しつつデコレータ化しました。その際に「デコレータ化した時のテストってどう書いておくべきだっけ」と考えて試してみた結果となります。

haoyayoi

2021.01.26

この記事は公開されてから1年以上経過しています。情報が古い可能性がありますので、ご注意ください。

はじめに

Pythonベースで作成されたAWS Lambdaのハンドラリファクタリングをしていました。個人的にとにかく面倒に感じたのはeventとcontextのバリデーションでした。値のチェックによる行数の肥大化、及び一時変数の増加がつきものになります。

チェック用コードの関数化、及び関数のdecorator化とまで進めてみましたが、中々頭を捻ったのがdecoretorのテストです。実際のコード及びテストを元にした実装例を交えて、考慮した部分をまとめてみました。

decoretorのコード

テスト対象のdecoratorコードは以下の通り。実際のコードから最低限のロジックを抜粋しています。decorator化にはpartialを使う手もありましたが、wrapsはラップされた関数のみを引数に指定すれば良い点がメリットです。

import json
from argparse import ArgumentParser
import datetime
from functools import wraps


MESSAGE_INVALID_DATE = '年月の形式が異なります'


def validate_date(yyyymm):
    try:
        datetime.datetime.strptime(f"{yyyymm}01", "%Y%m01")
        return yyyymm
    except ValueError:
        raise argparse.ArgumentTypeError(f"{MESSAGE_INVALID_DATE} - {yyyymm}")


def validate_id(id):
    if not bool(re.search(r'^[0-9]{12}$', id)):
        return
    return id


def parse_arguments(func):
    @wraps(func)
    def wrapper(event, context, *args, **kwds):
        validate_data = []
        for record in event['Records']:
            info = json.loads(record['body'])
            cmd_params = info['command_params']

            parser = ArgumentParser()
            parser.add_argument('yyyymm', help='年月', type=validate_date)
            parser.add_argument('acc_ids', help='アカウントID',
                                type=validate_id, nargs='+')
            data = parser.parse_args(cmd_params)
            validate_data.append({
                'yyyymm': data.yyyymm,
                'acc_ids': data.acc_ids
            })

        kwds['validate_data'] = validate_data
        return func(event, context, *args, **kwds)
    return wrapper

実際にdecoratorとして使う場合は、以下のようになります。

from ... import parse_arguments

@parse_arguments
def invoke(*args, **kwargs):
    ....

decoratorのテストを追加する

decoratorとして使った状態でのテストを追加していきます。wrapperとしてdecoratorが素のパラメータを受け取り、チェックし、wrappされた関数に渡します。本来のハンドラと異なり、あくまでもvalidate後のデータ比較のため、今回の関数は何もせずに直接返す構成にしました。

import pytest
import json


@pytest.mark.parametrize('args, result', [
    ('202101', True),
    ('202121', False)
])
def test_validate_date(args, result):
    if result:
        assert validate_date(args) == args
    else:
        with pytest.raises(argparse.ArgumentTypeError):
            validate_date(args)


@pytest.mark.parametrize('args, result, valid_id', [
    ('123456789012', True, '123456789012'),
    ('123456789', False, None),
    ('1234567890123', False, None),
])
def test_validate_ids(args, result, valid_id):
    assert validate_id(args) == valid_id


@pytest.mark.parametrize('arg, context, params', [
    ({
        "Records": [{
            "body": json.dumps({
                "command_params": [
                    "202101", "123456789023", "123456789012", "123456789023"
                ]
            })
        }]
    }, {}, [{
        'yyyymm': '202101',
        'acc_ids': ["123456789023", "123456789012", "123456789023"],
    }])
])
def test_parse_arguments(arg, context, params):

    @parse_arguments
    def test(*args, **kwargs):
        return kwargs['validate_data']

    assert test(arg, context) == params

parse_argumentsそのものに対してのアサーションも考えられますが、decorator動作そのものはPythonライブラリ側で担保されているものと考え、今回は省きました。

テストの実行

テスト実施後の結果です。一部加工しています。

% cat pytest.ini
[pytest]
log_cli = 1
log_cli_level = ERROR
testpaths = ./tests/ ./lambda_function/
python_files = test_*.py
python_paths = ./lambda_function/

% pipenv run pytest --flake8 --cache-clear -vv --cov=lambda_function

Courtesy Notice: Pipenv found itself running within a virtual environment, so it will automatically use that environment, instead of creating its own for any project. You can set PIPENV_IGNORE_VIRTUALENVS=1 to force pipenv to ignore that environment and create its own instead. You can set PIPENV_VERBOSITY=-1 to suppress this warning.
================================================================== test session starts ====================================================================
platform darwin -- Python 3.7.5, pytest-6.2.1, py-1.10.0, pluggy-0.13.1 -- .venv/bin/python
cachedir: .pytest_cache
rootdir: ., configfile: pytest.ini, testpaths: ./tests/, ./lambda_function/
plugins: flake8-1.0.7, env-0.6.2, cov-2.11.1, pythonpath-0.7.3
collected 42 items

tests/test_decorator.py::FLAKE8 PASSED
tests/test_decorator.py::test_validate_date[202101-True] PASSED                                                                                                                    
tests/test_decorator.py::test_validate_date[202121-False] PASSED                                            
tests/test_decorator.py::test_validate_ids[123456789012-True-123456789012] PASSED                                                                                                      
tests/test_decorator.py::test_validate_ids[123456789-False-None] PASSED                                                                                                                 
tests/test_decorator.py::test_validate_ids[1234567890123-False-None] PASSED
tests/test_decorator.py::test_parse_arguments[arg0-context0-params0] PASSED

---------- coverage: platform darwin, python 3.7.5-final-0 -----------
Name                                               Stmts   Miss  Cover
----------------------------------------------------------------------
lambda_function/__init__.py                            0      0   100%
lambda_function/task.py                               38      0   100%
----------------------------------------------------------------------
TOTAL                                                 38      0   100%

================================================================== 1 passed in 0.43s ======================================================================

あとがき

AWS Lambdaの引数がdict型の入れ子となっている場合や、値のチェックパターンが多岐に及ぶ場合、decorator化した関数で扱うメリットは大いにあります。逆にシンプルだとハンドラ内でのチェックで済ませるほうがメンテナンスコストも低く済みます。

コードそのものに対して求める精度にもよりますが、dict型の入れ子が連なったLambda関数に手を焼いている場合は検討してみることをおすすめします。

AWS