1 files changed, 43 insertions, 39 deletions

diff --git a/spec/lib/gitlab/query_limiting_spec.rb b/spec/lib/gitlab/query_limiting_spec.rb
index 4f70c65adca..fbb12629056 100644
--- a/spec/lib/gitlab/query_limiting_spec.rb
+++ b/spec/lib/gitlab/query_limiting_spec.rb
@@ -2,81 +2,85 @@
 
 require 'spec_helper'
 
-RSpec.describe Gitlab::QueryLimiting do
-  describe '.enable?' do
+RSpec.describe Gitlab::QueryLimiting, :request_store do
+  describe '.enabled_for_env?' do
     it 'returns true in a test environment' do
-      expect(described_class.enable?).to eq(true)
+      expect(described_class.enabled_for_env?).to eq(true)
     end
 
     it 'returns true in a development environment' do
       stub_rails_env('development')
       stub_rails_env('development')
 
-      expect(described_class.enable?).to eq(true)
+      expect(described_class.enabled_for_env?).to eq(true)
     end
 
     it 'returns false on GitLab.com' do
       stub_rails_env('production')
       allow(Gitlab).to receive(:com?).and_return(true)
 
-      expect(described_class.enable?).to eq(false)
+      expect(described_class.enabled_for_env?).to eq(false)
     end
 
     it 'returns false in a non GitLab.com' do
       allow(Gitlab).to receive(:com?).and_return(false)
       stub_rails_env('production')
 
-      expect(described_class.enable?).to eq(false)
+      expect(described_class.enabled_for_env?).to eq(false)
     end
   end
 
-  describe '.whitelist' do
-    it 'raises ArgumentError when an invalid issue URL is given' do
-      expect { described_class.whitelist('foo') }
-        .to raise_error(ArgumentError)
+  shared_context 'disable and enable' do |result|
+    let(:transaction) { Gitlab::QueryLimiting::Transaction.new }
+    let(:code) do
+      proc do
+        2.times { User.count }
+      end
     end
 
-    context 'without a transaction' do
-      it 'does nothing' do
-        expect { described_class.whitelist('https://example.com') }
-          .not_to raise_error
-      end
+    before do
+      allow(Gitlab::QueryLimiting::Transaction)
+        .to receive(:current)
+        .and_return(transaction)
     end
+  end
 
-    context 'with a transaction' do
-      let(:transaction) { Gitlab::QueryLimiting::Transaction.new }
+  describe '.disable!' do
+    include_context 'disable and enable'
 
-      before do
-        allow(Gitlab::QueryLimiting::Transaction)
-          .to receive(:current)
-          .and_return(transaction)
-      end
+    it 'raises an ArgumentError when an invalid issue URL is given' do
+      expect { described_class.disable!('foo') }
+        .to raise_error(ArgumentError)
+    end
 
-      it 'does not increment the number of SQL queries executed in the block' do
-        before = transaction.count
+    it 'stops the number of SQL queries from being incremented' do
+      described_class.disable!('https://example.com')
 
-        described_class.whitelist('https://example.com')
+      expect { code.call }.not_to change { transaction.count }
+    end
+  end
 
-        2.times do
-          User.count
-        end
+  describe '.enable!' do
+    include_context 'disable and enable'
 
-        expect(transaction.count).to eq(before)
-      end
+    it 'allows the number of SQL queries to be incremented' do
+      described_class.enable!
 
-      it 'whitelists when enabled' do
-        described_class.whitelist('https://example.com')
+      expect { code.call }.to change { transaction.count }.by(2)
+    end
+  end
 
-        expect(transaction.whitelisted).to eq(true)
-      end
+  describe '#enabled?' do
+    it 'returns true when enabled' do
+      Gitlab::SafeRequestStore[:query_limiting_disabled] = nil
 
-      it 'does not whitelist when disabled' do
-        allow(described_class).to receive(:enable?).and_return(false)
+      expect(described_class).to be_enabled
+    end
 
-        described_class.whitelist('https://example.com')
+    it 'returns false when disabled' do
+      Gitlab::SafeRequestStore[:query_limiting_disabled] = true
 
-        expect(transaction.whitelisted).to eq(false)
-      end
+      expect(described_class).not_to be_enabled
     end
   end
 end