Python
Flask
pytest

FlaskでAPI制限を実装しテストをする

More than 1 year has passed since last update.

FlaskでAPI Limitを実装しようとして色々情報収集に時間を費やしてしまったので放出

FlaskでAPI制限を利用するには、本体の機能だけではかなり面倒。そこでFlask-Limiterを用います。


インストール

以下のコマンドでインストールできます。


terminal

$ pip install Flask-Limiter



実装


app.py

# -*- coding: utf-8 -*-

from flask import Flask
from flask_limiter import Limiter
from flask_limiter.util import get_remote_address

app = Flask(__name__)
limiter = Limiter(
app,
key_func=get_remote_address,
global_limits=['100 per day', '10 per hour']
)

@app.route('/')
@limiter.limit('10/day')
def root():
return '1日10回しかアクセス出来ないぞ'

def main():
app.run()

if __name__ == '__main__':
main()


ちょっと長いですが、順を追って説明していきます。

まず limiter のインスタンスを作ります。

この時、引数にアプリケーション本体とkey_funcには関数を。また、global_limits (任意)を渡します。

key_func にここでは get_remote_address を渡しています。この関数はリクエストを送ってきたIPアドレスを返す関数です。ここでは、何を基準に規制かどうか判断する関数を渡してあげればいいです。

limiter.limit は実際にAPI制限をつけるデコレータです。記法は複数種類に対応しており、;でつなげたり、多重にデコレートすることができます。

区切り文字
記述例
意味

per
1 per second
1秒に1リクエストまで

/
1/day
1日1リクエストまで

引数を指定しない場合は、limiter コンストラクタに渡した global_limits の値が参照されます。


テストの仕方

本記事ではpytestを用いてテストコードを書いたり実行したりします。


terminal

$ pip install pytest


また、先のサンプルだと余計なエンドポイントが多いのでここで一旦整理と、テストのための準備をします。


app.py

# -*- coding: utf-8 -*-

from flask import Flask
from flask_limiter import Limiter
from flask_limiter.util import get_remote_address

app = Flask(__name__)
app.config['API_LIMIT'] = '50 per day'
limiter = Limiter(
app,
key_func=get_remote_address,
)

@app.route('/')
@limiter.limit(app.config['API_LIMIT'])
def root():
return '1日10回しかアクセス出来ないぞ'

def main():
app.run()

if __name__ == '__main__':
main()


ポイントとしては、テスト環境のときだけAPI制限がかかるまでのアクセス数をものすごく小さい数にしたいので app.config にAPI制限周りのパラメータを格納しました。

ポイントを考慮してテストコードを書くと以下の通りになります。


test_app.py

# -*- coding: utf-8 -*-

from app import app

class TestApplication:
@classmethod
def setup_class(cls):
cls.client = app.test_client()

def test_invalid(self):
tmp = app.config['API_LIMIT']
app.config['API_LIMIT'] = '1 per day'
response = self.client.get('/')
assert response.status_code == 200
response = self.client.get('/')
assert response.status_code == 429
app.config['API_LIMIT'] = tmp


テストを実行してみましょう!

py.test test_app.py

テストは落ちます。

================================== FAILURES ==================================

________________________ TestApplication.test_invalid ________________________

self = <test_app.TestApplication object at 0x10ed32860>

def test_invalid(self):
tmp = app.config["API_LIMIT"]
app.config["API_LIMIT"] = '1 per day'
response = self.client.get('/')
assert response.status_code == 200
response = self.client.get('/')
> assert response.status_code == 429
E assert 200 == 429
E + where 200 = <Response streamed [200 OK]>.status_code

test_app.py:15: AssertionError
========================== 1 failed in 0.26 seconds ==========================

これを解決するには、 limier.limitに渡す引数をLIMIT数を表す文字列から、LIMIT数を表す文字列を返す関数に変えてあげればOKです。:ok_woman:

app.py を書き換えてあげましょう。


app.py

# -*- coding: utf-8 -*-

from flask import Flask
from flask_limiter import Limiter
from flask_limiter.util import get_remote_address

app = Flask(__name__)
app.config['API_LIMIT'] = '50 per day'
limiter = Limiter(
app,
key_func=get_remote_address,
)

def limit_string():
return app.config['API_LIMIT']

@app.route('/')
@limiter.limit(limit_string)
def root():
return '1日10回しかアクセス出来ないぞ'

def main():
app.run()

if __name__ == '__main__':
main()