Merge branch 'bvl-extend-query-recorder' into 'master'

Extend the QueryRecorder matcher

See merge request gitlab-org/gitlab-ce!14267

3 files changed, 39 insertions, 7 deletions

diff --git a/spec/features/milestones/show_spec.rb b/spec/features/milestones/show_spec.rb
index 624f13922ed..50c5e0bb65f 100644
--- a/spec/features/milestones/show_spec.rb
+++ b/spec/features/milestones/show_spec.rb
@@ -18,9 +18,9 @@ describe 'Milestone show' do
 
   it 'avoids N+1 database queries' do
     create(:labeled_issue, issue_params)
-    control_count = ActiveRecord::QueryRecorder.new { visit_milestone }.count
+    control = ActiveRecord::QueryRecorder.new { visit_milestone }
     create_list(:labeled_issue, 10, issue_params)
 
-    expect { visit_milestone }.not_to exceed_query_limit(control_count)
+    expect { visit_milestone }.not_to exceed_query_limit(control)
   end
 end
diff --git a/spec/requests/api/projects_spec.rb b/spec/requests/api/projects_spec.rb
index 9602584f546..92e7d797cbd 100644
--- a/spec/requests/api/projects_spec.rb
+++ b/spec/requests/api/projects_spec.rb
@@ -54,9 +54,9 @@ describe API::Projects do
 
     shared_examples_for 'projects response without N + 1 queries' do
       it 'avoids N + 1 queries' do
-        control_count = ActiveRecord::QueryRecorder.new do
+        control = ActiveRecord::QueryRecorder.new do
           get api('/projects', current_user)
-        end.count
+        end
 
         if defined?(additional_project)
           additional_project
@@ -66,7 +66,7 @@ describe API::Projects do
 
         expect do
           get api('/projects', current_user)
-        end.not_to exceed_query_limit(control_count + 8)
+        end.not_to exceed_query_limit(control).with_threshold(8)
       end
     end
 
diff --git a/spec/support/query_recorder.rb b/spec/support/query_recorder.rb
index 55b531b4cf7..ba0b805caad 100644
--- a/spec/support/query_recorder.rb
+++ b/spec/support/query_recorder.rb
@@ -34,15 +34,47 @@ RSpec::Matchers.define :exceed_query_limit do |expected|
   supports_block_expectations
 
   match do |block|
-    query_count(&block) > expected
+    query_count(&block) > expected_count + threshold
   end
 
   failure_message_when_negated do |actual|
-    "Expected a maximum of #{expected} queries, got #{@recorder.count}:\n\n#{@recorder.log_message}"
+    threshold_message = threshold > 0 ? " (+#{@threshold})" : ''
+    counts = "#{expected_count}#{threshold_message}"
+    "Expected a maximum of #{counts} queries, got #{actual_count}:\n\n#{log_message}"
+  end
+
+  def with_threshold(threshold)
+    @threshold = threshold
+    self
+  end
+
+  def threshold
+    @threshold.to_i
+  end
+
+  def expected_count
+    if expected.is_a?(ActiveRecord::QueryRecorder)
+      expected.count
+    else
+      expected
+    end
+  end
+
+  def actual_count
+    @recorder.count
   end
 
   def query_count(&block)
     @recorder = ActiveRecord::QueryRecorder.new(&block)
     @recorder.count
   end
+
+  def log_message
+    if expected.is_a?(ActiveRecord::QueryRecorder)
+      extra_queries = (expected.log - @recorder.log).join("\n\n")
+      "Extra queries: \n\n #{extra_queries}"
+    else
+      @recorder.log_message
+    end
+  end
 end