Count occurrences of a specific query in the query recorder.

1 files changed, 14 insertions, 5 deletions

diff --git a/spec/support/query_recorder.rb b/spec/support/query_recorder.rb
index 369775db462..8cf8f45a8b2 100644
--- a/spec/support/query_recorder.rb
+++ b/spec/support/query_recorder.rb
@@ -41,7 +41,8 @@ RSpec::Matchers.define :exceed_query_limit do |expected|
   supports_block_expectations
 
   match do |block|
-    query_count(&block) > expected_count + threshold
+    @subject_block = block
+    actual_count > expected_count + threshold
   end
 
   failure_message_when_negated do |actual|
@@ -55,6 +56,11 @@ RSpec::Matchers.define :exceed_query_limit do |expected|
     self
   end
 
+  def for_query(query)
+    @query = query
+    self
+  end
+
   def threshold
     @threshold.to_i
   end
@@ -68,12 +74,15 @@ RSpec::Matchers.define :exceed_query_limit do |expected|
   end
 
   def actual_count
-    @recorder.count
+    @actual_count ||= if @query
+                        recorder.log.select { |recorded| recorded =~ @query }.size
+                      else
+                        recorder.count
+                      end
   end
 
-  def query_count(&block)
-    @recorder = ActiveRecord::QueryRecorder.new(&block)
-    @recorder.count
+  def recorder
+    @recorder ||= ActiveRecord::QueryRecorder.new(&@subject_block)
   end
 
   def count_queries(queries)