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)