Flask のテストクライアントをカスタムしてデフォルトのHTTPヘッダを設定する

Pythonの軽量フレームワークFlaskで自動テストをするためのテストクライアントをカスタムして、リクエスト毎にデフォルトのHTTPヘッダを設定する方法をご紹介します
2018.08.22

Flask のテストクライアントをカスタムしてデフォルトのHTTPヘッダを設定する

西田@大阪です

Pythonの軽量フレームワークFlask では自動テストをするためのテストクライアントが用意されています

今回はそのテストクライアントをカスタムして、リクエスト毎にデフォルトのHTTPヘッダを設定する方法をご紹介します

試した環境

  • Python 3.6.2
  • Flask 1.0.2

やりたいこと

今回はREST APIのテストをします

登場するAPIは以下の2つです

  • /authenticate 認証を行いaccess_tokenを返します
  • /products 商品情報を返します。認証が必要です

/authenticate のレスポンスに含まれる access_tokenAuthorizationヘッダに含めると/productsにアクセスすることができます。

/auhenticate

リクエスト

{
    "username": "user01",
    "password": "xxxxxx"
}

レスポンス

{
    "access_token": "xxx.xxxx.xxx"
}

/products

Authorization HTTPヘッダに/authenticateのレスポンスに含まれるaccess_tokenを設定することで認証できます

ヘッダ

Authorization: Bearer xxx.xxx.xxx

レスポンス

[
    {
        "id": 123,
        "name": "サンプル商品01"
    },
    {
        "id": 124,
        "name": "サンプル商品02"
    }
]

実装

カスタムしたtest_clientのクラスを作成します

FlaskClient を継承して作成します

from flask.testing import FlaskClient

class SecuredFlaskClient(FlaskClient):
    def __init__(self, *args, **kwargs):
        # 拡張したパラメーターをそのまま FlaskClient の
        # コンストラクタに渡さないために pop でとりだします
        access_token = kwargs.pop('access_token')
        super(SecuredFlaskClient, self).__init__(*args, **kwargs)

        # FlaskClient.envron_baseに'HTTP_' から始まるキーで値を設定すると、
        # Flask(WSGI)に渡すデフォルトのリクエストヘッダーになることを利用して
        # 初期化時に渡された access_token をリクエスト毎に 
        # Authorization リクエストヘッダに追加します
        self.environ_base['HTTP_AUTHORIZATION'] = f"Bearer {access_token}"

test_client を返すメソッドを用意します

from flaskr import flaskr

# デフォルトのテストクライアントを返します。認証のいらないAPIをテストするのに使います
def test_client():
    # test_client のクラスをデフォルトのFlaskClientに設定する
    flaskr.app.test_client_class = None
    return flaskr.app.test_client()

# カスタムされたテストクライアントを返します。認証の必要なAPIをテストするのに使います
def secured_test_client(username, password):
    # 認証を行うAPIを呼び出します
    rv = test_client().post('/authenticate', json=dict(username=username, password=password))

    # test_client のクラスを作成した SecuredFlaskClient に設定する
    flaskr.app.test_client_class = SecuredFlaskClient

    # テストクライアントに認証APIのレスポンスの `access_token` を
    # 渡してインスタンスを生成します
    return flaskr.app.test_client(access_token=rv.json['access_token'])

さいごに

いかがでしたでしょうか?

もちろん、テストクライアントをカスタムせずに、すべてのテストで認証、認証の必要なAPIの呼び出し、と順番に行えば同じことが実現できますが、コードが冗長になってしまいます

テストクライアントカスタムすれば、認証処理を一箇所にまとめられ変更に強くなり、記述するコード量を減らせておすすめです