From 3738358135dba55d120053f111cf6258419746fd Mon Sep 17 00:00:00 2001
From: Michael Elfassy <michael.elfassy@shopify.com>
Date: Mon, 9 Jan 2017 15:40:32 -0500
Subject: [PATCH] assert_enqueued_jobs with queue option

---
 activejob/lib/active_job/test_helper.rb  | 34 +++++++++++++++++-------
 activejob/test/cases/test_helper_test.rb | 21 +++++++++++++++
 2 files changed, 45 insertions(+), 10 deletions(-)

diff --git a/activejob/lib/active_job/test_helper.rb b/activejob/lib/active_job/test_helper.rb
index d01795f0c5833..7d621d6f578bd 100644
--- a/activejob/lib/active_job/test_helper.rb
+++ b/activejob/lib/active_job/test_helper.rb
@@ -55,7 +55,7 @@ def queue_adapter_for_test
     #     assert_enqueued_jobs 2
     #   end
     #
-    # If a block is passed, that block should cause the specified number of
+    # If a block is passed, that block will cause the specified number of
     # jobs to be enqueued.
     #
     #   def test_jobs_again
@@ -77,14 +77,23 @@ def queue_adapter_for_test
     #       HelloJob.perform_later('jeremy')
     #     end
     #   end
-    def assert_enqueued_jobs(number, only: nil)
+    #
+    # The number of times a job is enqueued to a specific queue can also be asserted.
+    #
+    #   def test_logging_job
+    #     assert_enqueued_jobs 2, queue: 'default' do
+    #       LoggingJob.perform_later
+    #       HelloJob.perform_later('elfassy')
+    #     end
+    #   end
+    def assert_enqueued_jobs(number, only: nil, queue: nil)
       if block_given?
-        original_count = enqueued_jobs_size(only: only)
+        original_count = enqueued_jobs_size(only: only, queue: queue)
         yield
-        new_count = enqueued_jobs_size(only: only)
+        new_count = enqueued_jobs_size(only: only, queue: queue)
         assert_equal number, new_count - original_count, "#{number} jobs expected, but #{new_count - original_count} were enqueued"
       else
-        actual_count = enqueued_jobs_size(only: only)
+        actual_count = enqueued_jobs_size(only: only, queue: queue)
         assert_equal number, actual_count, "#{number} jobs expected, but #{actual_count} were enqueued"
       end
     end
@@ -323,11 +332,16 @@ def clear_performed_jobs
         performed_jobs.clear
       end
 
-      def enqueued_jobs_size(only: nil)
-        if only
-          enqueued_jobs.count { |job| Array(only).include?(job.fetch(:job)) }
-        else
-          enqueued_jobs.count
+      def enqueued_jobs_size(only: nil, queue: nil)
+        enqueued_jobs.count do |job|
+          job_class = job.fetch(:job)
+          if only
+            next false unless Array(only).include?(job_class)
+          end
+          if queue
+            next false unless queue.to_s == job.fetch(:queue, job_class.queue_name)
+          end
+          true
         end
       end
 
diff --git a/activejob/test/cases/test_helper_test.rb b/activejob/test/cases/test_helper_test.rb
index 51fc6cc0c4552..5488ce3d5817b 100644
--- a/activejob/test/cases/test_helper_test.rb
+++ b/activejob/test/cases/test_helper_test.rb
@@ -110,6 +110,27 @@ def test_assert_enqueued_jobs_with_only_option
     end
   end
 
+  def test_assert_enqueued_jobs_with_only_and_queue_option
+    assert_nothing_raised do
+      assert_enqueued_jobs 1, only: HelloJob, queue: :some_queue do
+        HelloJob.set(queue: :some_queue).perform_later
+        HelloJob.set(queue: :other_queue).perform_later
+        LoggingJob.perform_later
+      end
+    end
+  end
+
+  def test_assert_enqueued_jobs_with_queue_option
+    assert_nothing_raised do
+      assert_enqueued_jobs 2, queue: :default do
+        HelloJob.perform_later
+        LoggingJob.perform_later
+        HelloJob.set(queue: :other_queue).perform_later
+        LoggingJob.set(queue: :other_queue).perform_later
+      end
+    end
+  end
+
   def test_assert_enqueued_jobs_with_only_option_and_none_sent
     error = assert_raise ActiveSupport::TestCase::Assertion do
       assert_enqueued_jobs 1, only: HelloJob do