def test_invalid_decorated_dynamic_limits_blueprint(self):
app = Flask(__name__)
app.config.setdefault("X", "2 per sec")
limiter = Limiter(app,
default_limits=["1/second"],
key_func=get_remote_address)
mock_handler = mock.Mock()
mock_handler.level = logging.INFO
limiter.logger.addHandler(mock_handler)
bp = Blueprint("bp1", __name__)
@bp.route("/t1")
def t1():
return "42"
limiter.limit(lambda: current_app.config.get("X"))(bp)
app.register_blueprint(bp)
with app.test_client() as cli:
with hiro.Timeline().freeze() as timeline:
self.assertEqual(cli.get("/t1").status_code, 200)
self.assertEqual(cli.get("/t1").status_code, 429)
self.assertEqual(mock_handler.handle.call_count, 3)
self.assertTrue("failed to load ratelimit" in
mock_handler.handle.call_args_list[0][0][0].msg)
self.assertTrue("failed to load ratelimit" in
mock_handler.handle.call_args_list[1][0][0].msg)
self.assertTrue("exceeded at endpoint" in
mock_handler.handle.call_args_list[2][0][0].msg)
def test_multiple_apps(self):
app1 = Flask(__name__)
app2 = Flask(__name__)
limiter = Limiter(default_limits=["1/second"],
key_func=get_remote_address)
limiter.init_app(app1)
limiter.init_app(app2)
@app1.route("/ping")
def ping():
return "PONG"
@app1.route("/slowping")
@limiter.limit("1/minute")
def slow_ping():
return "PONG"
@app2.route("/ping")
@limiter.limit("2/second")
def ping_2():
return "PONG"
@app2.route("/slowping")
@limiter.limit("2/minute")
def slow_ping_2():
return "PONG"
with hiro.Timeline().freeze() as timeline:
with app1.test_client() as cli:
self.assertEqual(cli.get("/ping").status_code, 200)
self.assertEqual(cli.get("/ping").status_code, 429)
timeline.forward(1)
self.assertEqual(cli.get("/ping").status_code, 200)
self.assertEqual(cli.get("/slowping").status_code, 200)
timeline.forward(59)
self.assertEqual(cli.get("/slowping").status_code, 429)
timeline.forward(1)
self.assertEqual(cli.get("/slowping").status_code, 200)
with app2.test_client() as cli:
self.assertEqual(cli.get("/ping").status_code, 200)
self.assertEqual(cli.get("/ping").status_code, 200)
self.assertEqual(cli.get("/ping").status_code, 429)
timeline.forward(1)
self.assertEqual(cli.get("/ping").status_code, 200)
self.assertEqual(cli.get("/slowping").status_code, 200)
timeline.forward(59)
self.assertEqual(cli.get("/slowping").status_code, 200)
self.assertEqual(cli.get("/slowping").status_code, 429)
timeline.forward(1)
self.assertEqual(cli.get("/slowping").status_code, 200)
def test_custom_headers_from_config(self):
app = Flask(__name__)
app.config.setdefault(C.HEADER_LIMIT, "X-Limit")
app.config.setdefault(C.HEADER_REMAINING, "X-Remaining")
app.config.setdefault(C.HEADER_RESET, "X-Reset")
limiter = Limiter(app,
default_limits=["10/minute"],
headers_enabled=True,
key_func=get_remote_address)
@app.route("/t1")
@limiter.limit("2/second; 10 per minute; 20/hour")
def t():
return "test"
with hiro.Timeline().freeze() as timeline:
with app.test_client() as cli:
for i in range(11):
resp = cli.get("/t1")
timeline.forward(1)
self.assertEqual(resp.headers.get('X-Limit'), '10')
self.assertEqual(resp.headers.get('X-Remaining'), '0')
self.assertEqual(resp.headers.get('X-Reset'),
str(int(time.time() + 50)))
def test_custom_headers_from_setter(self):
app = Flask(__name__)
limiter = Limiter(app,
default_limits=["10/minute"],
headers_enabled=True,
key_func=get_remote_address,
retry_after='http-date')
limiter._header_mapping[HEADERS.RESET] = 'X-Reset'
limiter._header_mapping[HEADERS.LIMIT] = 'X-Limit'
limiter._header_mapping[HEADERS.REMAINING] = 'X-Remaining'
@app.route("/t1")
@limiter.limit("2/second; 10 per minute; 20/hour")
def t():
return "test"
with hiro.Timeline().freeze(0) as timeline:
with app.test_client() as cli:
for i in range(11):
resp = cli.get("/t1")
timeline.forward(1)
self.assertEqual(resp.headers.get('X-Limit'), '10')
self.assertEqual(resp.headers.get('X-Remaining'), '0')
self.assertEqual(resp.headers.get('X-Reset'),
str(int(time.time() + 50)))
self.assertEqual(resp.headers.get('Retry-After'),
'Thu, 01 Jan 1970 00:01:01 GMT')
def test_headers_no_breach(self):
app = Flask(__name__)
limiter = Limiter(app,
default_limits=["10/minute"],
headers_enabled=True,
key_func=get_remote_address)
@app.route("/t1")
def t1():
return "test"
@app.route("/t2")
@limiter.limit("2/second; 5 per minute; 10/hour")
def t2():
return "test"
with hiro.Timeline().freeze():
with app.test_client() as cli:
resp = cli.get("/t1")
self.assertEqual(resp.headers.get('X-RateLimit-Limit'), '10')
self.assertEqual(resp.headers.get('X-RateLimit-Remaining'),
'9')
self.assertEqual(resp.headers.get('X-RateLimit-Reset'),
str(int(time.time() + 61)))
self.assertEqual(resp.headers.get('Retry-After'), str(60))
resp = cli.get("/t2")
self.assertEqual(resp.headers.get('X-RateLimit-Limit'), '2')
self.assertEqual(resp.headers.get('X-RateLimit-Remaining'),
'1')
self.assertEqual(resp.headers.get('X-RateLimit-Reset'),
str(int(time.time() + 2)))
self.assertEqual(resp.headers.get('Retry-After'), str(1))
def test_whitelisting(self):
app = Flask(__name__)
limiter = Limiter(app,
default_limits=["1/minute"],
headers_enabled=True,
key_func=get_remote_address)
@app.route("/")
def t():
return "test"
@limiter.request_filter
def w():
if request.headers.get("internal", None) == "true":
return True
return False
with hiro.Timeline().freeze() as timeline:
with app.test_client() as cli:
self.assertEqual(cli.get("/").status_code, 200)
self.assertEqual(cli.get("/").status_code, 429)
timeline.forward(60)
self.assertEqual(cli.get("/").status_code, 200)
for i in range(0, 10):
self.assertEqual(
cli.get("/", headers={
"internal": "true"
}).status_code, 200)
def test_conditional_limits(self):
"""Test that the conditional activation of the limits work."""
app = Flask(__name__)
limiter = Limiter(app, key_func=get_remote_address)
@app.route("/limited")
@limiter.limit("1 per day")
def limited_route():
return "passed"
@app.route("/unlimited")
@limiter.limit("1 per day", exempt_when=lambda: True)
def never_limited_route():
return "should always pass"
is_exempt = False
@app.route("/conditional")
@limiter.limit("1 per day", exempt_when=lambda: is_exempt)
def conditionally_limited_route():
return "conditional"
with app.test_client() as cli:
self.assertEqual(cli.get("/limited").status_code, 200)
self.assertEqual(cli.get("/limited").status_code, 429)
self.assertEqual(cli.get("/unlimited").status_code, 200)
self.assertEqual(cli.get("/unlimited").status_code, 200)
self.assertEqual(cli.get("/conditional").status_code, 200)
self.assertEqual(cli.get("/conditional").status_code, 429)
is_exempt = True
self.assertEqual(cli.get("/conditional").status_code, 200)
is_exempt = False
self.assertEqual(cli.get("/conditional").status_code, 429)
def test_constructor_arguments_over_config(redis_connection):
app = Flask(__name__)
app.config.setdefault(C.STRATEGY, "fixed-window-elastic-expiry")
limiter = Limiter(strategy='moving-window', key_func=get_remote_address)
limiter.init_app(app)
app.config.setdefault(C.STORAGE_URL, "redis://localhost:36379")
assert type(limiter._limiter) == MovingWindowRateLimiter
limiter = Limiter(storage_uri='memcached://localhost:31211',
key_func=get_remote_address)
limiter.init_app(app)
assert type(limiter._storage) == MemcachedStorage
def test_custom_headers_from_setter_and_config():
app = Flask(__name__)
app.config.setdefault(C.HEADER_LIMIT, "Limit")
app.config.setdefault(C.HEADER_REMAINING, "Remaining")
app.config.setdefault(C.HEADER_RESET, "Reset")
limiter = Limiter(default_limits=["10/minute"],
headers_enabled=True,
key_func=get_remote_address)
limiter._header_mapping[HEADERS.REMAINING] = 'Available'
limiter.init_app(app)
@app.route("/t1")
def t():
return "test"
with app.test_client() as cli:
for i in range(11):
resp = cli.get("/t1")
assert resp.headers.get('Limit') == '10'
assert resp.headers.get('Available') == '0'
assert resp.headers.get('Reset') is not None
def test_uninitialized_limiter():
app = Flask(__name__)
limiter = Limiter(default_limits=["1/hour"])
@app.route("/")
@limiter.limit("2/hour")
def root():
return "root"
with app.test_client() as client:
assert client.get("/").status_code == 200
assert client.get("/").status_code == 200
assert client.get("/").status_code == 200
def test_logging(caplog):
app = Flask(__name__)
limiter = Limiter(app, key_func=get_remote_address)
@app.route("/t1")
@limiter.limit("1/minute")
def t1():
return "test"
with app.test_client() as cli:
assert 200 == cli.get("/t1").status_code
assert 429 == cli.get("/t1").status_code
assert len(caplog.records) == 1
assert caplog.records[0].levelname == 'WARNING'
def test_logging(self):
app = Flask(__name__)
limiter = Limiter(app, key_func=get_remote_address)
mock_handler = mock.Mock()
mock_handler.level = logging.INFO
limiter.logger.addHandler(mock_handler)
@app.route("/t1")
@limiter.limit("1/minute")
def t1():
return "test"
with app.test_client() as cli:
self.assertEqual(200, cli.get("/t1").status_code)
self.assertEqual(429, cli.get("/t1").status_code)
self.assertEqual(mock_handler.handle.call_count, 1)
def test_retry_after_exists_rfc1123():
app = Flask(__name__)
_ = Limiter(app,
default_limits=["1/minute"],
headers_enabled=True,
key_func=get_remote_address)
@app.route("/t1")
def t():
return "", 200, {'Retry-After': 'Sun, 06 Nov 2032 01:01:01 GMT'}
with app.test_client() as cli:
resp = cli.get("/t1")
retry_after = int(resp.headers.get('Retry-After'))
assert retry_after > 1000
def test_retry_after(self):
app = Flask(__name__)
_ = Limiter(app,
default_limits=["1/minute"],
headers_enabled=True,
key_func=get_remote_address)
@app.route("/t1")
def t():
return "test"
with hiro.Timeline().freeze() as timeline:
with app.test_client() as cli:
resp = cli.get("/t1")
retry_after = int(resp.headers.get('Retry-After'))
self.assertTrue(retry_after > 0)
timeline.forward(retry_after)
resp = cli.get("/t1")
self.assertEqual(resp.status_code, 200)
def test_reuse_logging(self):
app = Flask(__name__)
app_handler = mock.Mock()
app_handler.level = logging.INFO
app.logger.addHandler(app_handler)
limiter = Limiter(app, key_func=get_remote_address)
for handler in app.logger.handlers:
limiter.logger.addHandler(handler)
@app.route("/t1")
@limiter.limit("1/minute")
def t1():
return "42"
with app.test_client() as cli:
cli.get("/t1")
cli.get("/t1")
self.assertEqual(app_handler.handle.call_count, 1)
def test_invalid_config_with_disabled():
app = Flask(__name__)
app.config.setdefault(C.ENABLED, False)
app.config.setdefault(C.STORAGE_URL, "fubar://")
limiter = Limiter(app, default_limits=["1/hour"])
@app.route("/")
def root():
return "root"
@app.route("/explicit")
@limiter.limit("2/hour")
def explicit():
return "explicit"
with app.test_client() as client:
assert client.get("/").status_code == 200
assert client.get("/").status_code == 200
assert client.get("/explicit").status_code == 200
assert client.get("/explicit").status_code == 200
assert client.get("/explicit").status_code == 200
def test_invalid_decorated_static_limits(self):
app = Flask(__name__)
limiter = Limiter(app,
default_limits=["1/second"],
key_func=get_remote_address)
mock_handler = mock.Mock()
mock_handler.level = logging.INFO
limiter.logger.addHandler(mock_handler)
@app.route("/t1")
@limiter.limit("2/sec")
def t1():
return "42"
with app.test_client() as cli:
with hiro.Timeline().freeze() as timeline:
self.assertEqual(cli.get("/t1").status_code, 200)
self.assertEqual(cli.get("/t1").status_code, 429)
self.assertTrue("failed to configure" in
mock_handler.handle.call_args_list[0][0][0].msg)
self.assertTrue("exceeded at endpoint" in
mock_handler.handle.call_args_list[1][0][0].msg)
def test_headers_breach():
app = Flask(__name__)
limiter = Limiter(app,
default_limits=["10/minute"],
headers_enabled=True,
key_func=get_remote_address)
@app.route("/t1")
@limiter.limit("2/second; 10 per minute; 20/hour")
def t():
return "test"
with hiro.Timeline().freeze() as timeline:
with app.test_client() as cli:
for i in range(11):
resp = cli.get("/t1")
timeline.forward(1)
assert resp.headers.get('X-RateLimit-Limit') == '10'
assert resp.headers.get('X-RateLimit-Remaining') == '0'
assert resp.headers.get('X-RateLimit-Reset') == \
str(int(time.time() + 50))
assert resp.headers.get('Retry-After') == str(int(50))
# init SeaSurf
seasurf = SeaSurf(app)
# init flask CDN
CDN(app)
# init flask HTMLMIN
HTMLMIN(app)
# init assets environment
assets = Environment(app)
register_filter(BabiliFilter)
class MiniJSONEncoder(JSONEncoder):
"""Minify JSON output."""
item_separator = ','
key_separator = ':'
def default(self, obj):
if isinstance(obj, datetime):
return obj.isoformat()+"Z"
# Let the base class default method raise the TypeError
return JSONEncoder.default(self, obj)
app.json_encoder = MiniJSONEncoder
# init rate limiting
limiter = Limiter(key_func=get_ipaddr, storage_uri="memory://", strategy="moving-window")
limiter.init_app(app)
def test_invalid_storage_string():
app = Flask(__name__)
app.config.setdefault(C.STORAGE_URL, "fubar://localhost:1234")
with pytest.raises(ConfigurationError):
Limiter(app, key_func=get_remote_address)
def test_invalid_strategy():
app = Flask(__name__)
app.config.setdefault(C.STRATEGY, "fubar")
with pytest.raises(ConfigurationError):
Limiter(app, key_func=get_remote_address)
