負荷試験ツールLocustで無限にダミーデータをPOSTする

新規案件で AWS Fargate を使ったアプリケーションを構築しています。負荷試験を行った際に、負荷試験ツール Locust でダミーデータを使った POST のテストをした話です。負荷試験なのでひたすらに負荷をかけていくのですが、1回 POST するときに1件ダミーデータを利用する場合、大量のダミーデータが必要になりますが、用意するのはとても大変です。ちゃんとやるなら本番を想定してあらゆる種類のダミーデータを用意すべきですが、今回は時間もないのと高パフォーマンスを求められるものではなかったため、少量のダミーデータで使い回すことにしました。本記事は Python 製負荷試験ツールで少量のダミーデータを使い回すための方法を紹介します。

Locust とは
サンプルコード1
サンプルコード2
最後に
参考URL

Locust とは

その前に Locust について説明します。Locust は Python 製の負荷試験ツールです。英語で「イナゴ、バッタ」を意味しており、対象に対して大量に飛び交うイメージをしてもらればと思います。Locust の特徴は以下のとおりです。

シナリオを Python で記述できるので、テストコードとしてソースコードに含めることができる
Web UI が用意されており、試験結果が容易に確認できる
- Web UI を使わずにコマンドラインから行うことも可能
試験結果がとてもシンプル（スループットとレイテンシ）なのでサッと試験したいときに便利
- Web UI クラスを拡張して独自の試験結果を表示することも可能
Docker Compose だとワーカー（攻撃するノード）を簡単に増やすことができて、高負荷をかけることができる

locust.io

今回の案件は Python で実装しており、シナリオスクリプトをソースコードで管理することで、他の人でも検証が容易にできると判断して Locust を採用しました。

サンプルコード1

/add_user に id, name, email を POST してユーザ登録する API を例にしました。ポイントは itertools.cycle() で JSON データを無限にイテレータしています。これにより3件しかないダミーデータでも無限に使い回すことができます。id は重複しないように毎回インクリメントするようにしています。

from locust import TaskSet, HttpUser, task, constant
import itertools

"""
ダミーデータ
idはPOSTされる度にインクリメントされる
"""
test_data = [{
    "id": "1",
    "name": "foo",
    "email": "foo@example.com"
}, {
    "id": "2",
    "name": "bar",
    "email": "bar@example.com"
}, {
    "id": "3",
    "name": "baz",
    "email": "baz@example.com"
}]

class MyMainTasks(TaskSet):
    count = 0

    def on_start(self):
        """
        テスト用データ取得
        itertools.cycle() で何度もデータを使い続けている
        """
        self.test_data = itertools.cycle(test_data)

    def add_user(self, data):
        self.client.post(url="/add_user",
                         headers={"Content-Type": "application/json"},
                         json=data)

    """
    シナリオの実行
    """
    @task(1)
    def scenario(self):
        data = next(self.test_data)
        data["id"] = self.count
        self.add_user(data)
        self.count += 1

class MyUser(HttpUser):
    wait_time = constant(1)  # 1秒おきにリクエストを送る
    tasks = [MyMainTasks]

サンプルコード2

今度はダミーデータを CSV ファイルに保存して読み込んで実行します。実行中は CSV ファイルをオープンにする必要があるので on_start() でオープンにし on_stop() でクローズします。

from locust import TaskSet, HttpUser, task, constant
import csv
import itertools

class MyMainTasks(TaskSet):
    count = 0

    def on_start(self):
        """
        テスト用データ取得
        itertools.cycle() で何度もデータを使い続けている
        """
        self.f = open('/path/to/test_data.csv', 'r')
        self.test_data = itertools.cycle(csv.reader(self.f))

    def on_stop(self):
        self.f.close()

    def add_user(self, data):
        self.client.post(url="/add_user",
                         headers={"Content-Type": "application/json"},
                         json=data)

    """
    シナリオの実行
    """
    @task(1)
    def scenario_ja(self):
        row = next(self.test_data)
        data = {}
        data["id"] = self.count
        data["name"] = row[1]
        data["email"] = row[2]
        self.add_user(data)
        self.count += 1

class MyUser(HttpUser):
    wait_time = constant(1)  # 1秒おきにリクエストを送る
    tasks = [MyMainTasks]