Python flask_limiter.Limiter() Examples

def test_invalid_decorated_dynamic_limits(caplog):
    app = Flask(__name__)
    app.config.setdefault("X", "2 per sec")
    limiter = Limiter(
        app, default_limits=["1/second"], key_func=get_remote_address
    )

    @app.route("/t1")
    @limiter.limit(lambda: current_app.config.get("X"))
    def t1():
        return "42"

    with app.test_client() as cli:
        with hiro.Timeline().freeze():
            assert cli.get("/t1").status_code == 200
            assert cli.get("/t1").status_code == 429
    # 2 for invalid limit, 1 for warning.
    assert len(caplog.records) == 3
    assert (
        "failed to load ratelimit"
        in caplog.records[0].msg
    )
    assert (
        "failed to load ratelimit"
        in caplog.records[1].msg
    )
    assert (
        "exceeded at endpoint"
        in caplog.records[2].msg
    )
    assert caplog.records[2].levelname == 'WARNING' 

def test_invalid_decorated_static_limits(caplog):
    app = Flask(__name__)
    limiter = Limiter(
        app, default_limits=["1/second"], key_func=get_remote_address
    )

    @app.route("/t1")
    @limiter.limit("2/sec")
    def t1():
        return "42"

    with app.test_client() as cli:
        with hiro.Timeline().freeze():
            assert cli.get("/t1").status_code == 200
            assert cli.get("/t1").status_code == 429
    assert (
        "failed to configure"
        in caplog.records[0].msg
    )
    assert (
        "exceeded at endpoint"
        in caplog.records[1].msg
    ) 

def test_invalid_decorated_dynamic_limits_blueprint(caplog):
    app = Flask(__name__)
    app.config.setdefault("X", "2 per sec")
    limiter = Limiter(
        app, default_limits=["1/second"], key_func=get_remote_address
    )
    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():
            assert cli.get("/t1").status_code == 200
            assert cli.get("/t1").status_code == 429

    assert len(caplog.records) == 3
    assert "failed to load ratelimit" in caplog.records[0].msg
    assert "failed to load ratelimit" in caplog.records[1].msg
    assert "exceeded at endpoint" in caplog.records[2].msg 

def __init__(self,global_limits=["1000/minute"]):
        self.app = Flask(__name__)
        self.global_limits = global_limits
        self.limiter = Limiter(self.app,key_func=get_ipaddr,global_limits=self.global_limits) 

def init(app):

    limiter = Limiter(
        app,
        key_func=get_remote_address,
        default_limits=[]
    )
    return limiter 

def test_conditional_limits():
    """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:
        assert cli.get("/limited").status_code == 200
        assert cli.get("/limited").status_code == 429

        assert cli.get("/unlimited").status_code == 200
        assert cli.get("/unlimited").status_code == 200

        assert cli.get("/conditional").status_code == 200
        assert cli.get("/conditional").status_code == 429
        is_exempt = True
        assert cli.get("/conditional").status_code == 200
        is_exempt = False
        assert cli.get("/conditional").status_code == 429 

def test_conditional_shared_limits():
    """Test that conditional shared limits work."""
    app = Flask(__name__)
    limiter = Limiter(app, key_func=get_remote_address)

    @app.route("/limited")
    @limiter.shared_limit("1 per day", "test_scope")
    def limited_route():
        return "passed"

    @app.route("/unlimited")
    @limiter.shared_limit(
        "1 per day", "test_scope", exempt_when=lambda: True
    )
    def never_limited_route():
        return "should always pass"

    is_exempt = False

    @app.route("/conditional")
    @limiter.shared_limit(
        "1 per day", "test_scope", exempt_when=lambda: is_exempt
    )
    def conditionally_limited_route():
        return "conditional"

    with app.test_client() as cli:
        assert cli.get("/unlimited").status_code == 200
        assert cli.get("/unlimited").status_code == 200

        assert cli.get("/limited").status_code == 200
        assert cli.get("/limited").status_code == 429

        assert cli.get("/conditional").status_code == 429
        is_exempt = True
        assert cli.get("/conditional").status_code == 200
        is_exempt = False
        assert cli.get("/conditional").status_code == 429 

def test_whitelisting():

    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:
            assert cli.get("/").status_code == 200
            assert cli.get("/").status_code == 429
            timeline.forward(60)
            assert cli.get("/").status_code == 200

            for i in range(0, 10):
                assert cli.get(
                    "/", headers={"internal": "true"}
                ).status_code == 200 

def test_invalid_decorated_static_limit_blueprint(caplog):
    app = Flask(__name__)
    limiter = Limiter(
        app, default_limits=["1/second"], key_func=get_remote_address
    )
    bp = Blueprint("bp1", __name__)

    @bp.route("/t1")
    def t1():
        return "42"

    limiter.limit("2/sec")(bp)
    app.register_blueprint(bp)

    with app.test_client() as cli:
        with hiro.Timeline().freeze():
            assert cli.get("/t1").status_code == 200
            assert cli.get("/t1").status_code == 429
    assert (
        "failed to configure"
        in caplog.records[0].msg
    )
    assert (
        "exceeded at endpoint"
        in caplog.records[1].msg
    ) 

def wrap(app, rule):
    """Wrap flask app with rule based request rate control.
    """
    limit_by = rule.pop('_limit_by', None)
    limiter = flask_limiter.Limiter(
        app,
        key_func=_get_key_func(limit_by=limit_by),
    )

    limit_value = rule.pop('_global', None)
    if limit_value is not None:
        decorator = limiter.shared_limit(limit_value, scope='_global')
        for endpoint, func in app.view_functions.items():
            # A shared rate limit could be evaluated multiple times for single
            # request, because `flask_limiter` uses
            # "func.__module__ + func.__name__" as key format to differentiate
            # non Blueprint routes.
            #
            # For example, both cases blow register to the same key format as
            # "flask.helpers.send_static_file"
            # 1. `restful_plus` static file requests (i.e. "/swaggerui/ui.js").
            # 2. default static file endpoint registration of
            # `treadmill.rest.FLASK_APP`, because param "static_folder" in
            # constructer has its default value.
            #
            # so each request of hitting "flask.helpers.send_static_file" will
            # increase rate limit counter by 2 (which 1 is expectation). A
            # workaround here is only concerning about `treadmill.rest.api`
            # resources.
            if hasattr(func, 'view_class'):
                app.view_functions[endpoint] = decorator(func)

    if rule:
        decorators = {
            module: limiter.shared_limit(limit_value, scope=module)
            for module, limit_value in rule.items()
        }

        for endpoint, func in app.view_functions.items():

            module = None
            if hasattr(func, 'view_class'):
                module = func.__module__.rsplit('.')[-1]

            if module in decorators:
                decorator = decorators[module]
                app.view_functions[endpoint] = decorator(func)

    return app