1 files changed, 70 insertions, 22 deletions

diff --git a/spec/support/matchers/exceed_query_limit.rb b/spec/support/matchers/exceed_query_limit.rb
index b48c7f905b2..bfcaf9552b3 100644
--- a/spec/support/matchers/exceed_query_limit.rb
+++ b/spec/support/matchers/exceed_query_limit.rb
@@ -1,6 +1,68 @@
 # frozen_string_literal: true
 
 module ExceedQueryLimitHelpers
+  class QueryDiff
+    def initialize(expected, actual, show_common_queries)
+      @expected = expected
+      @actual = actual
+      @show_common_queries = show_common_queries
+    end
+
+    def diff
+      return combined_counts if @show_common_queries
+
+      combined_counts
+        .transform_values { select_suffixes_with_diffs(_1) }
+        .reject { |_prefix, suffs| suffs.empty? }
+    end
+
+    private
+
+    def select_suffixes_with_diffs(suffs)
+      reject_groups_with_different_parameters(reject_suffixes_with_identical_counts(suffs))
+    end
+
+    def reject_suffixes_with_identical_counts(suffs)
+      suffs.reject { |_k, counts| counts.first == counts.second }
+    end
+
+    # Eliminates groups that differ only in parameters,
+    # to make it easier to debug the output.
+    #
+    # For example, if we have a group `SELECT * FROM users...`,
+    # with the following suffixes
+    #      `WHERE id = 1` (counts: N, 0)
+    #      `WHERE id = 2` (counts: 0, N)
+    def reject_groups_with_different_parameters(suffs)
+      return suffs if suffs.size != 2
+
+      counts_a, counts_b = suffs.values
+      return {} if counts_a == counts_b.reverse && counts_a.include?(0)
+
+      suffs
+    end
+
+    def expected_counts
+      @expected.transform_values do |suffixes|
+        suffixes.transform_values { |n| [n, 0] }
+      end
+    end
+
+    def recorded_counts
+      @actual.transform_values do |suffixes|
+        suffixes.transform_values { |n| [0, n] }
+      end
+    end
+
+    def combined_counts
+      expected_counts.merge(recorded_counts) do |_k, exp, got|
+        exp.merge(got) do |_k, exp_counts, got_counts|
+          exp_counts.zip(got_counts).map { |a, b| a + b }
+        end
+      end
+    end
+  end
+
   MARGINALIA_ANNOTATION_REGEX = %r{\s*\/\*.*\*\/}.freeze
 
   DB_QUERY_RE = Regexp.union([
@@ -108,26 +170,7 @@ module ExceedQueryLimitHelpers
   end
 
   def diff_query_counts(expected, actual)
-    expected_counts = expected.transform_values do |suffixes|
-      suffixes.transform_values { |n| [n, 0] }
-    end
-    recorded_counts = actual.transform_values do |suffixes|
-      suffixes.transform_values { |n| [0, n] }
-    end
-
-    combined_counts = expected_counts.merge(recorded_counts) do |_k, exp, got|
-      exp.merge(got) do |_k, exp_counts, got_counts|
-        exp_counts.zip(got_counts).map { |a, b| a + b }
-      end
-    end
-
-    unless @show_common_queries
-      combined_counts = combined_counts.transform_values do |suffs|
-        suffs.reject { |_k, counts| counts.first == counts.second }
-      end
-    end
-
-    combined_counts.reject { |_prefix, suffs| suffs.empty? }
+    QueryDiff.new(expected, actual, @show_common_queries).diff
   end
 
   def diff_query_group_message(query, suffixes)
@@ -141,7 +184,7 @@ module ExceedQueryLimitHelpers
   def log_message
     if expected.is_a?(ActiveRecord::QueryRecorder)
       diff_counts = diff_query_counts(count_queries(expected), count_queries(@recorder))
-      sections = diff_counts.map { |q, suffixes| diff_query_group_message(q, suffixes) }
+      sections = diff_counts.filter_map { |q, suffixes| diff_query_group_message(q, suffixes) }
 
       <<~MSG
       Query Diff:
@@ -323,7 +366,12 @@ RSpec::Matchers.define :exceed_query_limit do |expected|
   include ExceedQueryLimitHelpers
 
   match do |block|
-    verify_count(&block)
+    if block.is_a?(ActiveRecord::QueryRecorder)
+      @recorder = block
+      verify_count
+    else
+      verify_count(&block)
+    end
   end
 
   failure_message_when_negated do |actual|