サマリ

任意のbotocore.response.StreamingBodyオブジェクトを作る場合

文字列をencode()でUTF-8のバイト列に変換
io.BytesIO()でbytesオブジェクトに変換
バイト列の長さと合わせてbotocore.response.StreamingBody()でオブジェクト生成。

本文

S3に置いたファイルをPython(boto3)で取得する時にget_objectを利用する以下の様なコードが題材。実行環境はlambdaでもローカルでも。

lambda_function.py

import os
# when you run it in local env.
# os.environ['AWS_ACCESS_KEY_ID'] = 'YOUR_ACCESS_KEY'
# os.environ['AWS_SECRET_ACCESS_KEY'] = 'YOUR_SECRET_ACCESS_KEY'
# os.environ['AWS_DEFAULT_REGION'] = 'YOUR_REGION'
import boto3
import json
import traceback

BUCKET_NAME = 'bucket'
OBJECT_KEY = 'hoge.json'
def lambda_handler(event, context):
    try:
        client = boto3.client('s3')
        res = client.get_object(
            Bucket=BUCKET_NAME,
            Key=OBJECT_KEY
        )
        return json.loads(res["Body"].read().decode("utf-8"))
    except:
        traceback.print_exc()
        raise Exception("500")

get_objectでは以下の様なレスポンスが来るが、ファイルの中身はbotocore.response.StreamingBody のオブジェクトになっている。
ここではUTにおいてオブジェクトの中身を任意の値に置き換えたオブジェクトを返すモックを用意したい場合の方法を記述する。

{'ResponseMetadata': {'RequestId': '〇〇', 'HostId': '〇〇', 'HTTPStatusCode': 200, 'HTTPHeaders': {'x-amz-id-2': '〇〇', 'x-amz-request-id': '〇〇', 'date': 'Sat, 18 Sep 2021 07:34:40 GMT', 'last-modified': 'Sat, 18 Sep 2021 05:44:58 GMT', 'etag': '"〇〇"', 'accept-ranges': 'bytes', 'content-type': 'application/json', 'server': 'AmazonS3', 'content-length': '16'}, 'RetryAttempts': 0}, 'AcceptRanges': 'bytes', 'LastModified': datetime.datetime(2021, 9, 18, 5, 44, 58, tzinfo=tzutc()), 'ContentLength': 16, 'ETag': '"〇〇"', 'ContentType': 'application/json', 'Metadata': {}, 'Body': }

準備

$ pip install boto3
$ pip install pytest pytest-mock

実行方法

$ pytest

正常系

test_lambda_function.py

import pytest
import datetime
import json
import boto3
import botocore
from botocore.stub import Stubber
from dateutil.tz import tzutc
from io import BytesIO
from botocore.exceptions import ClientError

from lambda_function import lambda_handler

# -- Nomal Case -----------------------------------------------
# testNomalCase1 : succeeded in getting file.
# -- Error Case -----------------------------------------------
# testErrorCase1 : failed to get file. (後述)
# testErrorCase2 : broken json format. (後述)

def testNomalCase1(mocker):
    client = boto3.client('s3')
    stubber = Stubber(client)
    
    # 返却したいオブジェクト
    body_json = {
        "aaa": 3,
        "bbb": [
            {
                "ccc": "ddd"
            }
        ]
    }
    # エンコード。(encode()はデフォルトでutf-8。)
    body_encoded = json.dumps(body_json).encode()
    # body_encoded = "{\"aaa\": 3}".encode() # 勿論直接dumps後の文字列を埋めるのでもok
    content_length = len(body_encoded)
    # StreamingBodyへ整形する。
    body = botocore.response.StreamingBody(
        BytesIO(body_encoded),
        content_length
    )

    # モックが返すレスポンス
    res = {
        "ResponseMetadata": {
            "RequestId": "MD3HBGSCABVXW9NN", 
            "HostId": "RbuSBmsNqGyTmB3iHRED9xMkLmWtdFLtm/U7oc...", 
            "HTTPStatusCode": 200, 
            "HTTPHeaders": {
                "x-amz-id-2": "RbuSBmsNqGyTmB3iHRED9xMkLmWtdFLtm/U7oc...",
                "x-amz-request-id": "MD3HBGZENBVX2PNN",
                "date": "Sat, 18 Sep 2021 09:57:30 GMT", 
                "last-modified": "Sat, 18 Sep 2021 05:44:58 GMT", 
                "etag": "\"9a61704303ce71c6b9e79933b099169c\"", 
                "accept-ranges": "bytes", 
                "content-type": "application/json", 
                "server": "AmazonS3", 
                "content-length": content_length
            }, 
            "RetryAttempts": 0
        }, 
        "AcceptRanges": "bytes", 
        "LastModified": datetime.datetime(2021, 9, 18, 5, 44, 58, tzinfo=tzutc()), 
        "ContentLength": content_length, 
        "ETag": "\"9a61704303ce71c6b9e79933b099169c\"", 
        "ContentType": "application/json", 
        "Metadata": {}, 
        "Body": body
    }
    
    # get_objectのレスポンスにセット。
    stubber.add_response('get_object', res)
    stubber.activate()
    mocker.patch('boto3.client', return_value=client)
    
    # モックの返す結果をreturnしているかをチェック。
    assert lambda_handler(event=None, context=None) == body_json

異常系

S3 APIが投げる例外を再現する

test_lambda_function.py

def testErrorCase1(mocker):
    client = boto3.client('s3')
    stubber = Stubber(client)

    stubber.add_client_error('get_object', service_error_code='NoSuchKey')
    stubber.activate()
    mocker.patch('boto3.client', return_value=client)

    # 例外を拾う時はwith句    
    with pytest.raises(Exception):
        _ = lambda_handler(event=None, context=None)

取得したjsonのloadsに失敗

test_lambda_function.py

def testErrorCase2(mocker):
    client = boto3.client('s3')
    stubber = Stubber(client)
    
    # jsonとして成立しない文字列をエンコード。
    body_encoded = "{\"aaa\": }".encode()
    content_length = len(body_encoded)
    body = botocore.response.StreamingBody(
        BytesIO(body_encoded),
        content_length
    )

    # モックが返すレスポンス
    res = {
        "ResponseMetadata": {
            "RequestId": "MD3HBGSCABVXW9NN", 
            "HostId": "RbuSBmsNqGyTmB3iHRED9xMkLmWtdFLtm/U7oc...", 
            "HTTPStatusCode": 200, 
            "HTTPHeaders": {
                "x-amz-id-2": "RbuSBmsNqGyTmB3iHRED9xMkLmWtdFLtm/U7oc...",
                "x-amz-request-id": "MD3HBGZENBVX2PNN",
                "date": "Sat, 18 Sep 2021 09:57:30 GMT", 
                "last-modified": "Sat, 18 Sep 2021 05:44:58 GMT", 
                "etag": "\"9a61704303ce71c6b9e79933b099169c\"", 
                "accept-ranges": "bytes", 
                "content-type": "application/json", 
                "server": "AmazonS3", 
                "content-length": content_length
            }, 
            "RetryAttempts": 0
        }, 
        "AcceptRanges": "bytes", 
        "LastModified": datetime.datetime(2021, 9, 18, 5, 44, 58, tzinfo=tzutc()), 
        "ContentLength": content_length, 
        "ETag": "\"9a61704303ce71c6b9e79933b099169c\"", 
        "ContentType": "application/json", 
        "Metadata": {}, 
        "Body": body
    }
    
    stubber.add_response('get_object', res)
    stubber.activate()
    mocker.patch('boto3.client', return_value=client)
    
    with pytest.raises(Exception):
        _ = lambda_handler(event=None, context=None)

参考

pytest

boto3

公式ドキュメント