From 08b2cc4d950cceb80d6fbbf886ac33fd2e9737cf Mon Sep 17 00:00:00 2001 From: Gonzalo Rodriguez Date: Wed, 21 Mar 2018 17:10:27 -0300 Subject: [PATCH] Acceptance test throttling with a dynamic period --- spec/acceptance/throttling_spec.rb | 50 ++++++++++++++++++++++++++++++ 1 file changed, 50 insertions(+) diff --git a/spec/acceptance/throttling_spec.rb b/spec/acceptance/throttling_spec.rb index e7b2361..233016c 100644 --- a/spec/acceptance/throttling_spec.rb +++ b/spec/acceptance/throttling_spec.rb @@ -62,4 +62,54 @@ describe "#throttle" do 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.get_header("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