From 666dc3d894372471ffd451fc50d86727f921a81c Mon Sep 17 00:00:00 2001
From: Gonzalo Rodriguez <grzuy@users.noreply.github.com>
Date: Thu, 15 Mar 2018 15:24:22 -0300
Subject: [PATCH] Acceptance test ability to customize blocked/throttled
 responses (#298)

* Acceptance test ability to customize blocked/throttled responses

* Don't let customizations to blocklisted/throttled responses leak to other test cases
---
 .../customizing_blocked_response_spec.rb      | 22 +++++++++++++++
 .../customizing_throttled_response_spec.rb    | 28 +++++++++++++++++++
 spec/spec_helper.rb                           |  8 ++++++
 3 files changed, 58 insertions(+)
 create mode 100644 spec/acceptance/customizing_blocked_response_spec.rb
 create mode 100644 spec/acceptance/customizing_throttled_response_spec.rb

diff --git a/spec/acceptance/customizing_blocked_response_spec.rb b/spec/acceptance/customizing_blocked_response_spec.rb
new file mode 100644
index 0000000..7f3f307
--- /dev/null
+++ b/spec/acceptance/customizing_blocked_response_spec.rb
@@ -0,0 +1,22 @@
+require_relative "../spec_helper"
+
+describe "Customizing block responses" do
+  it "can be customized" do
+    Rack::Attack.blocklist("block 1.2.3.4") do |request|
+      request.ip == "1.2.3.4"
+    end
+
+    get "/", {}, "REMOTE_ADDR" => "1.2.3.4"
+
+    assert_equal 403, last_response.status
+
+    Rack::Attack.blocklisted_response = lambda do |env|
+      [503, {}, ["Blocked"]]
+    end
+
+    get "/", {}, "REMOTE_ADDR" => "1.2.3.4"
+
+    assert_equal 503, last_response.status
+    assert_equal "Blocked", last_response.body
+  end
+end
diff --git a/spec/acceptance/customizing_throttled_response_spec.rb b/spec/acceptance/customizing_throttled_response_spec.rb
new file mode 100644
index 0000000..7f66372
--- /dev/null
+++ b/spec/acceptance/customizing_throttled_response_spec.rb
@@ -0,0 +1,28 @@
+require_relative "../spec_helper"
+
+describe "Customizing throttled response" do
+  it "can be customized" do
+    Rack::Attack.cache.store = ActiveSupport::Cache::MemoryStore.new
+
+    Rack::Attack.throttle("by ip", limit: 1, 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
+
+    Rack::Attack.throttled_response = lambda do |env|
+      [503, {}, ["Throttled"]]
+    end
+
+    get "/", {}, "REMOTE_ADDR" => "1.2.3.4"
+
+    assert_equal 503, last_response.status
+    assert_equal "Throttled", last_response.body
+  end
+end
diff --git a/spec/spec_helper.rb b/spec/spec_helper.rb
index 7cfc1b5..3473eba 100644
--- a/spec/spec_helper.rb
+++ b/spec/spec_helper.rb
@@ -23,9 +23,17 @@ class MiniTest::Spec
 
   include Rack::Test::Methods
 
+  before do
+    @_original_throttled_response = Rack::Attack.throttled_response
+    @_original_blocklisted_response = Rack::Attack.blocklisted_response
+  end
+
   after do
     Rack::Attack.clear!
     Rack::Attack.instance_variable_set(:@cache, nil)
+
+    Rack::Attack.throttled_response = @_original_throttled_response
+    Rack::Attack.blocklisted_response = @_original_blocklisted_response
   end
 
   def app