mirror of
https://github.com/laurentS/slowapi.git
synced 2026-03-13 09:10:20 +08:00
42 lines
1.3 KiB
Python
42 lines
1.3 KiB
Python
import logging
|
|
import platform
|
|
import unittest
|
|
from functools import wraps
|
|
|
|
import redis
|
|
from fastapi import FastAPI
|
|
from mock import mock
|
|
from starlette.applications import Starlette
|
|
|
|
from slowapi.errors import RateLimitExceeded
|
|
from slowapi.extension import Limiter, _rate_limit_exceeded_handler
|
|
from slowapi.util import get_remote_address
|
|
from slowapi.middleware import SlowAPIMiddleware
|
|
|
|
|
|
class TestSlowapi:
|
|
def build_starlette_app(self, config={}, **limiter_args):
|
|
limiter_args.setdefault("key_func", get_remote_address)
|
|
limiter = Limiter(**limiter_args)
|
|
app = Starlette(debug=True)
|
|
app.state.limiter = limiter
|
|
app.add_exception_handler(RateLimitExceeded, _rate_limit_exceeded_handler)
|
|
app.add_middleware(SlowAPIMiddleware)
|
|
|
|
mock_handler = mock.Mock()
|
|
mock_handler.level = logging.INFO
|
|
limiter.logger.addHandler(mock_handler)
|
|
return app, limiter
|
|
|
|
def build_fastapi_app(self, config={}, **limiter_args):
|
|
limiter_args.setdefault("key_func", get_remote_address)
|
|
limiter = Limiter(**limiter_args)
|
|
app = FastAPI()
|
|
app.state.limiter = limiter
|
|
app.add_middleware(SlowAPIMiddleware)
|
|
|
|
mock_handler = mock.Mock()
|
|
mock_handler.level = logging.INFO
|
|
limiter.logger.addHandler(mock_handler)
|
|
return app, limiter
|