AWS SDKのモック実装例

# はじめに

AWS SDKを使用したコードに対して高速且つ多様なケースの単体テストを実装するためのノウハウ(≒SDKの挙動のモック化)を、実装例と共に備忘録として残します。

---

## なぜAWS SDKをモック化したいのか

AWS SDKを使用したコードをモック化しないで単体テストしようとすると、以下のような点が実現性や生産性のボトルネックになりがちです。

- 実際にAWSリソースを操作しようとすると、手間と時間とコストが掛かる
- 実際のAWSのAPIから意図したレスポンス(例えば特定状況下でのエラーレスポンス)を返却させるのが困難、又は不可能
- セキュリティ上の都合でIAMやNWが制限されておりAWSのAPIの利用が不可能

以上のような理由から、高速且つ多様なケースをテストするためにAWS SDKをモック化したい場合があります。

---

## 注意・免責事項

- この記事に記載している内容が唯一且つ最適解という訳では無いかも知れません。もっと効率的なやる方があるかも知れません
- モック化はあくまでも単体テストでのみ使用すべきであり、結合試験・E2E試験時には実際のAWSリソースを使用して試験すべきでしょう

---

## 例題

例えば以下のような処理を実装したいとしましょう。

> 指定したIDのEC2インスタンスのステータスを数秒間隔で確認する。
ステータスが`running`になったら処理を正常終了する。
最大試行回数を超えてもステータスが`running`にならなかったら処理を異常終了する。

例えばPythonでの素朴な(必要最小限の)実装例は以下のような感じになると思います。
このコードを、AWS SDK(DescribeInstancesの挙動)をモック化せず単体テストしようとすると、まず前処理としてテスト用のEC2インスタンスの作成処理と、後処理としてそのインスタンスの削除処理が必要になります。
また、実際のインスタンスのステータス遷移のタイミングを細かく制御することが困難なので、テスト結果が非決定的になります。

```python
import time
import boto3

def main():
    ec2_client = boto3.client("ec2")
    instance_id = "i-12345678"
    retry_count = 5
    interval_sec = 5

    for _ in range(retry_count):
        time.sleep(interval_sec)
        response = ec2_client.describe_instances(
            InstanceIds=[instance_id]
        )
        reservations = response.get("Reservations", [])
        assert len(reservations) == 1
        instances = reservations[0].get("Instances", [])
        assert len(instances) == 1
        state = instances[0].get("State", None)
        assert state is not None
        state_name = state.get("Name", None)
        assert state_name is not None
        
        # ステータスが`running`であればループを抜ける
        if state_name == "running":
            break

    else:
        raise Exception("最大試行回数を超えてもインスタンスのステータスが`running`になりませんでした")

```

上記のような、AWS SDKを使用するコードであっても高速に実行出来て、結果が決定的になる単体テストを実装できるよう、AWS SDKのモック化の手順を例示していきます。

---

## 実装例

ここでは筆者が使用経験のある、Python、TypeScript、Goの実装例を記載していきます。

---

### Python

必要な手順は大きく以下の通りです。

1. AWS SDKの必要な部分をプロトコルとして定義する
2. テスト対象の関数の引数としてそのプロトコルを設定する
3. そのプロトコルを満たすダミーのクラスを定義しテストで使用する

<details><summary>アプリケーションコード</summary>

```python:sample.py
from __future__ import annotations
from time import sleep
from abc import abstractmethod
import typing
from typing import Protocol
# Python3.10以前はtyping_extensionsから、3.11以降はtypingからインポート可能
from typing_extensions import Unpack

import boto3
# 型チェック時にのみインポートすることで、アプリケーションデプロイ時にmypy_boto3ライブラリをインストールする必要が無くなる
# ※Lambda関数のzipのサイズ節約や手間の省略に繋がる
if typing.TYPE_CHECKING:
    # ※boto3の型情報はmypy_boto3ライブラリをインストールすることで利用可能
    # https://pypi.org/project/mypy-boto3/
    from mypy_boto3_ec2.type_defs import (
        DescribeInstancesResultTypeDef, DescribeInstancesRequestRequestTypeDef
    )

# 1. AWS SDKの必要な部分をプロトコルとして定義する
class EC2ClientProtpcol(Protocol):
    @abstractmethod
    def describe_instances(
        self,
        **kwargs: Unpack[DescribeInstancesRequestRequestTypeDef]
    ) -> DescribeInstancesResultTypeDef:
        pass

# 2. テスト対象の関数の引数としてそのプロトコルを設定する
def wait_until_instance_ready(
    ec2_client:EC2ClientProtpcol, instance_id:str,
    retry_count:int, interval_sec:int,
):
    for _ in range(retry_count):
        sleep(interval_sec)
        response = ec2_client.describe_instances(InstanceIds=[instance_id])
        reservations = response.get("Reservations", [])
        assert len(reservations) == 1
        instances = reservations[0].get("Instances", [])
        assert len(instances) == 1
        state = instances[0].get("State", None)
        assert state is not None
        state_name = state.get("Name", None)
        assert state_name is not None

        if state_name == "running":
            break
    else:
        raise Exception("最大試行回数を超えてもインスタンスのステータスが`running`になりませんでした")

def main():
    ec2_client = boto3.client("ec2")
    instance_id = "i-12345678"
    retry_count = 5
    interval_sec = 5

    wait_until_instance_ready(
        ec2_client, instance_id,
        retry_count, interval_sec,
    )
```

</details>

<details><summary>単体テストコード</summary>

```python:test_sample.py

import pytest
from typing_extensions import Unpack
from sample import EC2ClientProtpcol, wait_until_instance_ready

from mypy_boto3_ec2.type_defs import (
    DescribeInstancesResultTypeDef,
    DescribeInstancesRequestRequestTypeDef
)
from mypy_boto3_ec2.literals import InstanceStateNameType

# 3. そのプロトコルを満たすダミーのクラスを定義しテストで使用する
class DummyEC2Client(EC2ClientProtpcol):
    def __init__(self) -> None:
        # `describe_instances`メソッドの呼び出し回数を初期化する
        self.__call_count = 0

    def describe_instances(
        self,
        **kwargs: Unpack[DescribeInstancesRequestRequestTypeDef]
    ) -> DescribeInstancesResultTypeDef:
        # このメソッドの呼び出し回数をインクリメントする
        self.__call_count+=1

        instance_ids = kwargs.get("InstanceIds", [])
        assert len(instance_ids) == 1

        # 3回目の呼び出し時に`running`のステータスを返却する
        state:InstanceStateNameType = "pending"
        if self.__call_count == 3:
            state = "running"

        # 必要最小限のダミーデータをレスポンスする
        return {
            "Reservations": [
                {
                    "Instances": [
                        {
                            "InstanceId": instance_ids[0], 
                            "State": {"Name": state}
                        }
                    ]}
            ],
            "ResponseMetadata": {
                "HostId": "",
                "HTTPHeaders": {},
                "HTTPStatusCode": 200,
                "RequestId": "",
                "RetryAttempts": 0
            }
        }

def test_wait_until_instance_ready_ok():
    """インスタンスのステータスを複数回確認した後に`running`状態になる"""
    wait_until_instance_ready(DummyEC2Client(), "i-12345678", retry_count=3, interval_sec=1)

def test_wait_until_instance_ready_timeout():
    """リトライ回数を超えてもインスタンスのステータスが`running`にならない場合は例外を上げる"""
    with pytest.raises(Exception) as ex:
        wait_until_instance_ready(DummyEC2Client(), instance_id="i-12345678", retry_count=2, interval_sec=1)
    assert str(ex.value) == "最大試行回数を超えてもインスタンスのステータスが`running`になりませんでした"

```

</details>

---

### TypeScript

必要な手順は大きく以下の通りです。

1. AWS SDKのクライアントクラスを継承したダミーのクラスとその振舞いを定義しテストで使用する

<details><summary>アプリケーションコード</summary>

```typescript:sample.ts

import { DescribeInstancesCommand, EC2Client } from "@aws-sdk/client-ec2"
import {setTimeout} from "timers/promises"

export const waitUntilInstanceReady = async (
    ec2Client:EC2Client, instanceId:string,
    retryCount:number, intervalSec:number
) => {
    let currentCount = 0
    while (currentCount < retryCount) {
        await setTimeout(intervalSec*1000)
        const response = await ec2Client.send(new DescribeInstancesCommand({
            InstanceIds: [instanceId]
        }))
        const reservations = response.Reservations ?? []
        if (reservations.length === 0) {throw Error()}
        const instances = reservations[0].Instances ?? []
        if (instances.length === 0) {throw Error()}
        const state = instances[0].State
        if (state === undefined) {throw Error()}
        const stateName = state.Name
        if (stateName === undefined) {throw Error()}

        if (stateName === "running") {
            break
        } else {
            currentCount++
        }
    }
    if (currentCount >= retryCount) {
        throw Error("最大試行回数を超えてもインスタンスのステータスが`running`になりませんでした")
    }
}

const main = async () => {
    const ec2Client = new EC2Client()
    const instanceId = "i-12345678"
    const retryCount = 5
    const intervalSec = 5

    await waitUntilInstanceReady(ec2Client, instanceId, retryCount, intervalSec)
}

```

</details>

<details><summary>単体テストコード</summary>

```typescript:sample.test.ts

import { 
    DescribeInstancesCommand, DescribeInstancesCommandOutput, EC2Client, 
    ServiceInputTypes, ServiceOutputTypes, EC2ClientConfig, 
    InstanceStateName
} from "@aws-sdk/client-ec2";
import { SmithyResolvedConfiguration } from "@smithy/smithy-client";
import { Command, HttpHandlerOptions, CheckOptionalClientConfig } from "@smithy/types";
import { waitUntilInstanceReady } from "./sample";

// 1. AWS SDKのクライアントクラスを継承したダミーのクラスとその振舞いを定義しテストで使用する
class DummyEC2Client extends EC2Client {
    private callCount: number

    // `DescribeInstances`の呼び出し回数を初期化する
    constructor(...[configuration]: CheckOptionalClientConfig<EC2ClientConfig>) {
        super()
        this.callCount = 0
    }
    // IDE(VSCode)で`send`と入力するとタブ補完で下記4行分のsendメソッドの定義が出力される。
    // ※正直この4行全ての記述が必要な理由は良く分かっていない、、
    send<InputType extends ServiceInputTypes, OutputType extends ServiceOutputTypes>(command: Command<ServiceInputTypes, InputType, ServiceOutputTypes, OutputType, SmithyResolvedConfiguration<HttpHandlerOptions>>, options?: HttpHandlerOptions | undefined): Promise<OutputType>;
    send<InputType extends ServiceInputTypes, OutputType extends ServiceOutputTypes>(command: Command<ServiceInputTypes, InputType, ServiceOutputTypes, OutputType, SmithyResolvedConfiguration<HttpHandlerOptions>>, cb: (err: any, data?: OutputType | undefined) => void): void;
    send<InputType extends ServiceInputTypes, OutputType extends ServiceOutputTypes>(command: Command<ServiceInputTypes, InputType, ServiceOutputTypes, OutputType, SmithyResolvedConfiguration<HttpHandlerOptions>>, options: HttpHandlerOptions, cb: (err: any, data?: OutputType | undefined) => void): void;
    // `DescribeInstances`コマンドの挙動を定義する
    send<InputType extends ServiceInputTypes, OutputType extends ServiceOutputTypes>(
        command: unknown, options?: unknown, cb?: unknown
    ): void | Promise<OutputType> {
        if (!(command instanceof DescribeInstancesCommand)) {
            throw Error()
        }
        // `DescribeInstances`の呼び出し回数をインクリメントする
        this.callCount++
        // 3回目の呼び出し回数時に`running`ステータスを返却する
        let state:InstanceStateName = "pending"
        if (this.callCount === 3) {
            state = "running"
        }
        const instanceIds = command.input.InstanceIds ?? []
        if (instanceIds.length === 0) {throw Error()}
        const response:DescribeInstancesCommandOutput = {
            "$metadata": {},
            "Reservations": [
                {
                    "Instances": [
                        {
                            "InstanceId": instanceIds[0],
                            "State": {
                                "Name": state
                            }
                        }
                    ]
                }
            ]
        }
        return Promise.resolve(response as OutputType)
    }
}

let ec2Client = new DummyEC2Client()
test("インスタンスのステータスを複数回確認した後に`running`状態になる", async () => {
    await waitUntilInstanceReady(ec2Client, "i-12345678", 3, 1)
})
test("リトライ回数を超えてもインスタンスのステータスが`running`にならない場合は例外を上げる", async () => {
    const f = async () => {
        await waitUntilInstanceReady(ec2Client, "i-12345678", 2, 1)
    }
    await expect(f).rejects.toThrow("最大試行回数を超えてもインスタンスのステータスが`running`になりませんでした")
}

```

</details>

---

### Go

必要な手順は大きく以下の通りです。

1. AWS SDKの必要な部分をインターフェイスとして定義する
2. テスト対象の関数の引数としてそのインターフェイスを設定する
3. そのインターフェイスを満たすダミーの構造体を定義しテストで使用する

<details><summary>アプリケーションコード</summary>

```go:sample.go

package main

import (
	"context"
	"errors"
	"time"

	"github.com/aws/aws-sdk-go-v2/config"
	"github.com/aws/aws-sdk-go-v2/service/ec2"
	"github.com/aws/aws-sdk-go-v2/service/ec2/types"
)

// 1. AWS SDKの必要な部分をインターフェイスとして定義する
type EC2ClientInterface interface {
	DescribeInstances(
		ctx context.Context,
		params *ec2.DescribeInstancesInput,
		optFns ...func(*ec2.Options),
	) (*ec2.DescribeInstancesOutput, error)
}

// 2. テスト対象の関数の引数としてそのインターフェイスを設定する
func WaitUntilInstanceReady(ec2Client EC2ClientInterface, instanceId string, ctx context.Context, retryCount, intervalSec int) error {
	for i := 0; i < retryCount; i++ {
		time.Sleep(time.Second * time.Duration(intervalSec))
		response, err := ec2Client.DescribeInstances(ctx, &ec2.DescribeInstancesInput{
			InstanceIds: []string{instanceId},
		})
		if err != nil {
			return err
		}
		reservations := response.Reservations
		if len(reservations) == 0 {
			return errors.New("")
		}
		instances := reservations[0].Instances
		if len(instances) == 0 {
			return errors.New("")
		}
		stateName := instances[0].State.Name
		if stateName == types.InstanceStateNameRunning {
			return nil
		}
	}
	return errors.New("最大試行回数を超えてもインスタンスのステータスが`running`になりませんでした")
}
func main() {
	cfg, err := config.LoadDefaultConfig(context.TODO())
	if err != nil {
		panic(err)
	}
	ec2Client := ec2.NewFromConfig(cfg)
	instanceId := "i-12345678"
	retryCount := 5
	intervalSec := 5

	WaitUntilInstanceReady(ec2Client, instanceId, context.TODO(), retryCount, intervalSec)

}


```

</details>

<details><summary>単体テストコード</summary>

```go:sample_test.go

package main

import (
	"context"
	"testing"

	"github.com/aws/aws-sdk-go-v2/service/ec2"
	"github.com/aws/aws-sdk-go-v2/service/ec2/types"
	"github.com/stretchr/testify/assert"
)

// 3. そのインターフェイスを満たすダミーの構造体を定義しテストで使用する
type DummyEC2Client struct {
	callCount int
}

func (c *DummyEC2Client) DescribeInstances(ctx context.Context, params *ec2.DescribeInstancesInput, optFns ...func(*ec2.Options)) (*ec2.DescribeInstancesOutput, error) {
	// このメソッドの呼び出し回数をインクリメントする
	c.callCount++

	var state types.InstanceStateName = types.InstanceStateNamePending
	instanceId := params.InstanceIds[0]

	// 3回目の呼び出し時に`running`のステータスを返却する
	if c.callCount == 3 {
		state = types.InstanceStateNameRunning
	}
	return &ec2.DescribeInstancesOutput{
		Reservations: []types.Reservation{{
			Instances: []types.Instance{{
				InstanceId: &instanceId, State: &types.InstanceState{Name: state},
			}},
		}},
	}, nil
}

// インスタンスのステータスを複数回確認した後に`running`状態になる
func TestWaitUntilInstanceReady_ok(t *testing.T) {
	ec2Client := &DummyEC2Client{callCount: 0}

	err := WaitUntilInstanceReady(ec2Client, "i-12345678", context.TODO(), 3, 1)
	assert.Nil(t, err)
}

// リトライ回数を超えてもインスタンスのステータスが`running`にならない場合は例外を上げる
func TestWaitUntilInstanceReady_timeout(t *testing.T) {
	ec2Client := &DummyEC2Client{callCount: 0}

	err := WaitUntilInstanceReady(ec2Client, "i-12345678", context.TODO(), 2, 1)
	assert.Equal(t, err.Error(), "最大試行回数を超えてもインスタンスのステータスが`running`になりませんでした")
}

```

</details>