Add spec for ActionRateLimiter

2 files changed, 29 insertions, 2 deletions

diff --git a/lib/gitlab/action_rate_limiter.rb b/lib/gitlab/action_rate_limiter.rb
index c3af583a3ed..7b231ac14ca 100644
--- a/lib/gitlab/action_rate_limiter.rb
+++ b/lib/gitlab/action_rate_limiter.rb
@@ -16,12 +16,12 @@ module Gitlab
       value = 0
 
       Gitlab::Redis::Cache.with do |redis|
-        cache_key = "action_rate_limiter:#{action}:#{key}"
+        cache_key = "action_rate_limiter:#{action.to_s}:#{key}"
         value = redis.incr(cache_key)
         redis.expire(cache_key, expiry_time) if value == 1
       end
 
-      value.to_i
+      value
     end
 
     def throttled?(key, threshold_value)
diff --git a/spec/lib/gitlab/action_rate_limiter_spec.rb b/spec/lib/gitlab/action_rate_limiter_spec.rb
new file mode 100644
index 00000000000..84661605623
--- /dev/null
+++ b/spec/lib/gitlab/action_rate_limiter_spec.rb
@@ -0,0 +1,27 @@
+require 'spec_helper'
+
+describe Gitlab::ActionRateLimiter do
+  let(:redis) { double('redis') }
+  let(:key) { 'user:1' }
+  let(:cache_key) { "action_rate_limiter:test_action:#{key}" }
+
+  subject { described_class.new(action: :test_action, expiry_time: 100) }
+
+  before do
+    allow(Gitlab::Redis::Cache).to receive(:with).and_yield(redis)
+  end
+
+  it 'increases the throttle count and sets the expire time' do
+    expect(redis).to receive(:incr).with(cache_key).and_return(1)
+    expect(redis).to receive(:expire).with(cache_key, 100)
+
+    expect(subject.throttled?(key, 1)).to be false
+  end
+
+  it 'returns true if the key is throttled' do
+    expect(redis).to receive(:incr).with(cache_key).and_return(2)
+    expect(redis).not_to receive(:expire)
+
+    expect(subject.throttled?(key, 1)).to be true
+  end
+end