Allow to reset state between tests

Author: fatkodima

lib/rack/attack.rb

@@ -14,6 +14,7 @@ class Attack
     class Error < StandardError; end
     class MisconfiguredStoreError < Error; end
     class MissingStoreError < Error; end
+    class IncompatibleStoreError < Error; end
 
     autoload :Check,                'rack/attack/check'
     autoload :Throttle,             'rack/attack/throttle'
@@ -53,6 +54,10 @@ def clear!
         @configuration.clear_configuration
       end
 
+      def reset!
+        cache.reset!
+      end
+
       extend Forwardable
       def_delegators(
         :@configuration,

lib/rack/attack/cache.rb

@@ -41,6 +41,17 @@ def delete(unprefixed_key)
         store.delete("#{prefix}:#{unprefixed_key}")
       end
 
+      def reset!
+        if store.respond_to?(:delete_matched)
+          store.delete_matched("#{prefix}*")
+        else
+          raise(
+            Rack::Attack::IncompatibleStoreError,
+            "Configured store #{store.class.name} doesn't respond to #delete_matched method"
+          )
+        end
+      end
+
       private
 
       def key_and_expiry(unprefixed_key, period)

lib/rack/attack/store_proxy/redis_proxy.rb

@@ -43,6 +43,19 @@ def delete(key, _options = {})
           rescuing { del(key) }
         end
 
+        def delete_matched(matcher, _options = nil)
+          cursor = "0"
+
+          rescuing do
+            # Fetch keys in batches using SCAN to avoid blocking the Redis server.
+            loop do
+              cursor, keys = scan(cursor, match: matcher, count: 1000)
+              del(*keys) unless keys.empty?
+              break if cursor == "0"
+            end
+          end
+        end
+
         private
 
         def rescuing

spec/rack_attack_spec.rb

@@ -99,4 +99,26 @@
       end
     end
   end
+
+  describe 'reset!' do
+    it 'raises an error when is not supported by cache store' do
+      Rack::Attack.cache.store = Class.new
+      assert_raises(Rack::Attack::IncompatibleStoreError) do
+        Rack::Attack.reset!
+      end
+    end
+
+    if defined?(Redis)
+      it 'should delete rack attack keys' do
+        redis = Redis.new
+        redis.set('key', 'value')
+        redis.set("#{Rack::Attack.cache.prefix}::key", 'value')
+        Rack::Attack.cache.store = redis
+        Rack::Attack.reset!
+
+        _(redis.get('key')).must_equal 'value'
+        _(redis.get("#{Rack::Attack.cache.prefix}::key")).must_be_nil
+      end
+    end
+  end
 end