diff --git a/spec/acceptance/throttling_spec.rb b/spec/acceptance/throttling_spec.rb index 2ad4605..396589a 100644 --- a/spec/acceptance/throttling_spec.rb +++ b/spec/acceptance/throttling_spec.rb @@ -2,9 +2,11 @@ require_relative "../spec_helper" require "timecop" describe "#throttle" do - it "allows one request per minute by IP" do + before do Rack::Attack.cache.store = ActiveSupport::Cache::MemoryStore.new + end + it "allows one request per minute by IP" do Rack::Attack.throttle("by ip", limit: 1, period: 60) do |request| request.ip end @@ -29,4 +31,85 @@ describe "#throttle" do assert_equal 200, last_response.status end end + + it "supports limit to be dynamic" do + # Could be used to have different rate limits for authorized + # vs general requests + limit_proc = lambda do |request| + if request.env["X-APIKey"] == "private-secret" + 2 + else + 1 + end + end + + Rack::Attack.throttle("by ip", limit: limit_proc, period: 60) do |request| + request.ip + end + + get "/", {}, "REMOTE_ADDR" => "1.2.3.4" + assert_equal 200, last_response.status + + get "/", {}, "REMOTE_ADDR" => "1.2.3.4" + assert_equal 429, last_response.status + + get "/", {}, "REMOTE_ADDR" => "5.6.7.8", "X-APIKey" => "private-secret" + assert_equal 200, last_response.status + + get "/", {}, "REMOTE_ADDR" => "5.6.7.8", "X-APIKey" => "private-secret" + assert_equal 200, last_response.status + + get "/", {}, "REMOTE_ADDR" => "5.6.7.8", "X-APIKey" => "private-secret" + assert_equal 429, last_response.status + end + + it "supports period to be dynamic" do + # Could be used to have different rate limits for authorized + # vs general requests + period_proc = lambda do |request| + if request.env["X-APIKey"] == "private-secret" + 10 + else + 30 + end + end + + Rack::Attack.throttle("by ip", limit: 1, period: period_proc) do |request| + request.ip + end + + # Using Time#at to align to start/end of periods exactly + # to achieve consistenty in different test runs + + Timecop.travel(Time.at(0)) do + get "/", {}, "REMOTE_ADDR" => "1.2.3.4" + assert_equal 200, last_response.status + + get "/", {}, "REMOTE_ADDR" => "1.2.3.4" + assert_equal 429, last_response.status + end + + Timecop.travel(Time.at(10)) do + get "/", {}, "REMOTE_ADDR" => "1.2.3.4" + assert_equal 429, last_response.status + end + + Timecop.travel(Time.at(30)) do + get "/", {}, "REMOTE_ADDR" => "1.2.3.4" + assert_equal 200, last_response.status + end + + Timecop.travel(Time.at(0)) do + get "/", {}, "REMOTE_ADDR" => "5.6.7.8", "X-APIKey" => "private-secret" + assert_equal 200, last_response.status + + get "/", {}, "REMOTE_ADDR" => "5.6.7.8", "X-APIKey" => "private-secret" + assert_equal 429, last_response.status + end + + Timecop.travel(Time.at(10)) do + get "/", {}, "REMOTE_ADDR" => "5.6.7.8", "X-APIKey" => "private-secret" + assert_equal 200, last_response.status + end + end end