🦆

外部APIを使ったPythonコードをテストする（pytest-mock, MagicMock）

2024/12/11に公開

Python

pytest

tech

外部APIを使ったPythonコードをテストしたい

テストコードを書いていたところ、外部のAPIを呼び出すコードに対してどのようにテストを作ればいいのかという問題に直面しました。例えば以下のようなコードです。

gen_haiku.py

import re
from openai import OpenAI

def gen_haiku(haiku_theme):
    """
    テーマに対して俳句を生成して返却する
    ひらがな以外が含まれる場合はエラーを返す
    """
    # OpenAi APIを呼び出す
    client = OpenAI()
    completion = client.chat.completions.create(
        model="gpt-4o",
        messages=[
            {
                "role": "system", 
                "content": f"""次のテーマに対して日本語のひらがなで俳句を生成してください。
                [生成例1]テーマ:ミカン, 出力:ふゆのひに みかんのかおり へやみたす
                [生成例2]テーマ:ふじのやま, 出力:ふじのやま ゆきにかがやく あさのそら"""
            },
            {
                "role": "user", 
                "content": f"テーマ:{haiku_theme}, 出力:"
            }
        ]
    )
    haiku = completion.choices[0].message.content

    # フォーマットエラーをチェックする（ひらがな以外のときにエラーにする）
    if not re.fullmatch(r"[\u3041-\u3096\s]+", haiku):
        raise ValueError("生成された俳句にひらがな以外の文字が含まれています: " + haiku)

    # 生成した俳句を返却する
    return haiku

このコードのgen_haiku関数は、OpenAIを呼び出してテーマに沿った俳句を生成します。

このgen_haiku関数のテストを考えたときに、テストを実行するたびに毎回OpenAI APIを呼び出していては利用料金がかさみますし、生成AIのため実行結果も変わってしまいます（＝テスト結果も変わってしまいます）。なので基本的に、テストを実行するたびに毎回外部のAPIを呼び出すわけにはいきません。

こういった場合は外部APIをモック化し、外部APIを呼び出すのではなく、外部APIのような振る舞いをするモックを呼び出すようにします。

モック化を行う手段はいろいろありますが、以前使ったpytest-mockとMagicMockが便利だったので使い方を残しておきます。

pytest-mockのインストール

pytest-mockはpipを用いて以下で簡単にインストールできます。

pip install pytest-mock

MagicMockのインストール

MagicMockはPython標準のソフトウェアテストライブラリunittestに入っているため、Pythonが使えれば特に追加のインストールは不要です。

実際に書いてみる

使い方を説明するために、先ほどのgen_haiku関数のサンプルコードに対してモックを使ったテストコードを書いてみます。以下のように書いてみました。

test_gen_haiku.py

import pytest
from unittest.mock import MagicMock
from src.gen_haiku import gen_haiku

def test_gen_haiku_俳句にひらがな以外が含まれていた場合(mocker):

    # Given: モックの用意
    mock_client = MagicMock()
    mock_client.chat.completions.create.return_value = MagicMock(
        choices=[MagicMock(message=MagicMock(content="とうきょうの ひかりさします タワーかな"))] # ひらがな以外が含まれるためエラーになるはず
    )
    mocker.patch("src.gen_haiku.OpenAI", return_value=mock_client)

    # When: gen_haiku()の実行
    with pytest.raises(ValueError) as e:
        gen_haiku("東京タワー")

    # Then: 結果の確認
    assert "生成された俳句にひらがな以外の文字が含まれています" in str(e.value)

こちらのテストでは、OpenAI APIをとうきょうのひかりさしますタワーかなを返すモックに置き換えています。gen_haiku関数は生成された俳句にひらがな以外が含まれている場合はエラーを返すため、モックにカタカナを含む俳句を返させて、正しくエラーを発生するかをテストしています。

細かいトピックについて解説してきます。

Given-When-Then

まずpytest-mockと関係ないところですが、上記のテストコードはGiven-When-Thenで書いています。
簡単にいうと、テストの前提条件は何で（Given）、どんな操作をすると（When）、どうなってほしいのか（Then）というのをわかりやすくブロック化したものになります。
可読性を向上させるために今回はGiven-When-Then構文で書いています。

mocker.patch

テストの引数にmockerを指定し、mocker.patchを使うことでモックに置き換えることができます。

mocker.patch("src.gen_haiku.OpenAI", return_value=mock_client)

上記ではsrc.gen_haiku.OpenAIをmocker.patchでモックに置き換えています。またreturn_valueでこのモックの返却値をmock_clientに指定しています。

MagicMock

OpenAI APIの返却値を再現するためにMagicMockを使っています。
MagicMockは以下のようにimportすることで使うことができます。

from unittest.mock import MagicMock

MagicMockは空のオブジェクトのようなもので、以下のようにすることで簡単に属性を追加することができます。この部分では、mock_clientに属性（chat.completions.create）を追加し、その属性の返却値をMagicMockを用いて定義しています。

mock_client = MagicMock()
mock_client.chat.completions.create.return_value = MagicMock(
    choices=[MagicMock(message=MagicMock(content="とうきょうの ひかりさします タワーかな"))] # ひらがな以外が含まれるためエラーになるはず
)

このとき、mock_client.chat.completions.createの返却値は実際にOpenAI APIが返す返却値と一致するようにしています。ただOpenAI APIの返却値をすべて模倣するのは大変であり、またその必要もないため、実際の返却値の一部だけをモック化しています。

haiku = completion.choices[0].message.content

OpenAIが生成したテキストは上記のように取得できるため、これを再現するべく、以下のようにモックの返却値を定義しています。

mock_client.chat.completions.create.return_value = MagicMock(
    choices=[MagicMock(message=MagicMock(content="とうきょうの ひかりさします タワーかな"))]

以上、pytest-mockのmocker.patchやMagicMockを使うことで外部APIをモックに置き換えて比較的簡単にテストができるようになるという話でした。