1 files changed, 60 insertions, 5 deletions
diff --git a/spec/support/query_recorder.rb b/spec/support/query_recorder.rb
index 55b531b4cf7..8cf8f45a8b2 100644
--- a/spec/support/query_recorder.rb
+++ b/spec/support/query_recorder.rb
@@ -8,7 +8,14 @@ module ActiveRecord
       ActiveSupport::Notifications.subscribed(method(:callback), 'sql.active_record', &block)
     end
 
+    def show_backtrace(values)
+      Rails.logger.debug("QueryRecorder SQL: #{values[:sql]}")
+      caller.each { |line| Rails.logger.debug("   --> #{line}") }
+    end
+
     def callback(name, start, finish, message_id, values)
+      show_backtrace(values) if ENV['QUERY_RECORDER_DEBUG']
+
       if values[:name]&.include?("CACHE")
         @cached << values[:sql]
       elsif !values[:name]&.include?("SCHEMA")
@@ -34,15 +41,63 @@ RSpec::Matchers.define :exceed_query_limit do |expected|
   supports_block_expectations
 
   match do |block|
-    query_count(&block) > expected
+    @subject_block = block
+    actual_count > 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 for_query(query)
+    @query = query
+    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
+    @actual_count ||= if @query
+                        recorder.log.select { |recorded| recorded =~ @query }.size
+                      else
+                        recorder.count
+                      end
+  end
+
+  def recorder
+    @recorder ||= ActiveRecord::QueryRecorder.new(&@subject_block)
+  end
+
+  def count_queries(queries)
+    queries.each_with_object(Hash.new(0)) { |query, counts| counts[query] += 1 }
   end
 
-  def query_count(&block)
-    @recorder = ActiveRecord::QueryRecorder.new(&block)
-    @recorder.count
+  def log_message
+    if expected.is_a?(ActiveRecord::QueryRecorder)
+      counts = count_queries(expected.log)
+      extra_queries = @recorder.log.reject { |query| counts[query] -= 1 unless counts[query].zero? }
+      extra_queries_display = count_queries(extra_queries).map { |query, count| "[#{count}] #{query}" }
+
+      (['Extra queries:'] + extra_queries_display).join("\n\n")
+    else
+      @recorder.log_message
+    end
   end
 end