1 files changed, 13 insertions, 3 deletions
diff --git a/spec/support/shared_examples/graphql/n_plus_one_query_examples.rb b/spec/support/shared_examples/graphql/n_plus_one_query_examples.rb
index faf1bb204c9..b4afde311ba 100644
--- a/spec/support/shared_examples/graphql/n_plus_one_query_examples.rb
+++ b/spec/support/shared_examples/graphql/n_plus_one_query_examples.rb
@@ -1,13 +1,23 @@
 # frozen_string_literal: true
-RSpec.shared_examples 'N+1 query check' do
+
+RSpec.shared_examples 'N+1 query check' do |threshold: 0, skip_cached: true|
   it 'prevents N+1 queries' do
     execute_query # "warm up" to prevent undeterministic counts
     expect(graphql_errors).to be_blank # Sanity check - ex falso quodlibet!
 
-    control = ActiveRecord::QueryRecorder.new { execute_query }
+    control = ActiveRecord::QueryRecorder.new(skip_cached: skip_cached) { execute_query }
     expect(control.count).to be > 0
 
     search_params[:iids] << extra_iid_for_second_query
-    expect { execute_query }.not_to exceed_query_limit(control)
+
+    expect { execute_query }.not_to exceed_query_count_limit(control, skip_cached: skip_cached, threshold: threshold)
+  end
+
+  def exceed_query_count_limit(control, skip_cached: true, threshold: 0)
+    if skip_cached
+      exceed_query_limit(control).with_threshold(threshold)
+    else
+      exceed_all_query_limit(control).with_threshold(threshold)
+    end
   end
 end