From 208530494e5d2c5c62a3e1c24489aae0e4935e3a Mon Sep 17 00:00:00 2001
From: Patricio Cano <suprnova32@gmail.com>
Date: Tue, 8 Nov 2016 16:31:37 -0600
Subject: Refactored initializer code to its own class and added tests

---
 config/initializers/sidekiq.rb            |  8 +-------
 lib/gitlab/sidekiq_throttler.rb           | 21 +++++++++++++++++++
 spec/lib/gitlab/sidekiq_throttler_spec.rb | 34 +++++++++++++++++++++++++++++++
 3 files changed, 56 insertions(+), 7 deletions(-)
 create mode 100644 lib/gitlab/sidekiq_throttler.rb
 create mode 100644 spec/lib/gitlab/sidekiq_throttler_spec.rb

diff --git a/config/initializers/sidekiq.rb b/config/initializers/sidekiq.rb
index 7cc5e396f98..b87b31d9697 100644
--- a/config/initializers/sidekiq.rb
+++ b/config/initializers/sidekiq.rb
@@ -29,13 +29,7 @@ Sidekiq.configure_server do |config|
   end
   Sidekiq::Cron::Job.load_from_hash! cron_jobs
 
-  if Gitlab::CurrentSettings.sidekiq_throttling_enabled?
-    factor = current_application_settings.sidekiq_throttling_factor
-
-    current_application_settings.sidekiq_throttling_queues.each do |queue|
-      Sidekiq::Queue[queue].limit = (factor * Sidekiq.options[:concurrency]).ceil
-    end
-  end
+  Gitlab::SidekiqThrottler.execute!
 
   # Database pool should be at least `sidekiq_concurrency` + 2
   # For more info, see: https://github.com/mperham/sidekiq/blob/master/4.0-Upgrade.md
diff --git a/lib/gitlab/sidekiq_throttler.rb b/lib/gitlab/sidekiq_throttler.rb
new file mode 100644
index 00000000000..771736e7606
--- /dev/null
+++ b/lib/gitlab/sidekiq_throttler.rb
@@ -0,0 +1,21 @@
+module Gitlab
+  class SidekiqThrottler
+    class << self
+      def execute!
+        if Gitlab::CurrentSettings.sidekiq_throttling_enabled?
+          current_application_settings.sidekiq_throttling_queues.each do |queue|
+            Sidekiq::Queue[queue].limit = set_limit
+          end
+        end
+      end
+
+      private
+
+      def set_limit
+        factor = current_application_settings.sidekiq_throttling_factor
+
+        (factor * Sidekiq.options[:concurrency]).ceil
+      end
+    end
+  end
+end
diff --git a/spec/lib/gitlab/sidekiq_throttler_spec.rb b/spec/lib/gitlab/sidekiq_throttler_spec.rb
new file mode 100644
index 00000000000..ac4a64c0f43
--- /dev/null
+++ b/spec/lib/gitlab/sidekiq_throttler_spec.rb
@@ -0,0 +1,34 @@
+require 'spec_helper'
+
+describe Gitlab::SidekiqThrottler do
+  before do
+    Sidekiq.options[:concurrency] = 35
+
+    stub_application_setting(
+      sidekiq_throttling_enabled: true,
+      sidekiq_throttling_factor: 0.1,
+      sidekiq_throttling_queues: %w[build project_cache]
+    )
+  end
+
+  describe '#set_limit' do
+    it 'returns the correct limit' do
+      expect(Gitlab::SidekiqThrottler.send(:set_limit)).to eq 4
+    end
+  end
+
+  describe '#execute!' do
+    it 'sets limits on the selected queues' do
+      Gitlab::SidekiqThrottler.execute!
+
+      expect(Sidekiq::Queue['build'].limit).to eq 4
+      expect(Sidekiq::Queue['project_cache'].limit).to eq 4
+    end
+
+    it 'does not set limits on other queues' do
+      Gitlab::SidekiqThrottler.execute!
+
+      expect(Sidekiq::Queue['merge'].limit).to be_nil
+    end
+  end
+end
-- 
cgit v1.2.1