【pytest】関数をモックしてテストする

はじめに

pytestでモックを使ったコードを実装していたので簡単にまとめてみます。

テストしたいこと(一部)

• fetch_data関数からのステータスコードが200のとき、正しいデータが返ってくること
• データにはfetch_data関数からのレスポンスと、table.get_item関数からのレスポンスが含まれる

テストする上でのポイント

• モックする関数が2つ(fetch_data、table.get_item)
• unittest.mockのpatchとMagicMockを使って関数の返り値をモック
• boto3がprofileを参照しにいくため、テストでもダミーのAWS認証情報が必要です(未設定だとおこられます)

テスト対象のコード

app/sample/handler.py

def main(event, context):
    r = fetch_data(event['authorizationToken'])
    status = r.status_code
    if status == 200:
        data = json.loads(r.text)
        context = {
            id: '12345'
        }
    else:
        context = {}
    table_response = table.get_item(Key={'id': 'hoge'}, ProjectionExpression="column")
    table_item = table_response.get('Item')
    response = {}
    if table_item and table_item.get('column'):
        response['config'] = 
        {
            hogehoe
        }
    response['context'] = context
    return response

テストのコード

app/tests/test_handler.py

import pytest
import os
from unittest.mock import patch, MagicMock
import json

from sample.handler import main


@pytest.fixture()
def fetch_data_response_200():
    response = MagicMock()
    response.status_code = 200
    response.text = json.dumps({
        'id': 'user123'
    })
    return response


@pytest.fixture()
def example_event():
    return {
        'authorizationToken': 'hogeToken',
        'methodArn': 'hogeArn'
    }


def test_main_200(fetch_data_response_200, example_event):
    with patch('sample.handler.fetch_data', return_value=fetch_data_response_200), \
         patch('sample.handler.table.get_item') as mock_get_item:
            mock_get_item.return_value = {
                'Item': {
                    'id': {'S': '123'},
                    'column': {'L': [{'S': 'hogehoge'}]}
                }
            }
            result = main(example_event, {})
            assert result['context']['id'] == '12345'

おわりに

今回はチームメンバーが実装したコードに後からテストを加えたのですが、関数毎に実装してくれていたのでテストが書きやすかったです。mainの中身を関数の呼び出しのみに留めることができればもっと書きやすくなりそうです。