Merge pull request #34339 from Edouard-chin/ec-mail-test-helper

Fix ActionMailer assertion not working for mail defining delivery_job:

7 files changed, 171 insertions, 9 deletions

diff --git a/actionmailer/CHANGELOG.md b/actionmailer/CHANGELOG.md
index 35574a2c73..ec6f249974 100644
--- a/actionmailer/CHANGELOG.md
+++ b/actionmailer/CHANGELOG.md
@@ -1,3 +1,8 @@
+*   Fix ActionMailer assertions not working when a Mail defines
+    a custom delivery job class
+
+    *Edouard Chin*
+
 *   Mails with multipart `format` blocks with implicit render now also check for
     a template name in options hash instead of only using the action name.
 
diff --git a/actionmailer/lib/action_mailer/test_helper.rb b/actionmailer/lib/action_mailer/test_helper.rb
index a4751916af..b3403e8331 100644
--- a/actionmailer/lib/action_mailer/test_helper.rb
+++ b/actionmailer/lib/action_mailer/test_helper.rb
@@ -34,7 +34,7 @@ module ActionMailer
     def assert_emails(number, &block)
       if block_given?
         original_count = ActionMailer::Base.deliveries.size
-        perform_enqueued_jobs(only: [ActionMailer::DeliveryJob, ActionMailer::Parameterized::DeliveryJob], &block)
+        perform_enqueued_jobs(only: ->(job) { delivery_job_filter(job) }, &block)
         new_count = ActionMailer::Base.deliveries.size
         assert_equal number, new_count - original_count, "#{number} emails expected, but #{new_count - original_count} were sent"
       else
@@ -90,7 +90,7 @@ module ActionMailer
     #     end
     #   end
     def assert_enqueued_emails(number, &block)
-      assert_enqueued_jobs number, only: [ ActionMailer::DeliveryJob, ActionMailer::Parameterized::DeliveryJob ], &block
+      assert_enqueued_jobs(number, only: ->(job) { delivery_job_filter(job) }, &block)
     end
 
     # Asserts that a specific email has been enqueued, optionally
@@ -125,10 +125,10 @@ module ActionMailer
     #   end
     def assert_enqueued_email_with(mailer, method, args: nil, queue: "mailers", &block)
       if args.is_a? Hash
-        job = ActionMailer::Parameterized::DeliveryJob
+        job = mailer.parameterized_delivery_job
         args = [mailer.to_s, method.to_s, "deliver_now", args]
       else
-        job = ActionMailer::DeliveryJob
+        job = mailer.delivery_job
         args = [mailer.to_s, method.to_s, "deliver_now", *args]
       end
 
@@ -151,7 +151,16 @@ module ActionMailer
     #     end
     #   end
     def assert_no_enqueued_emails(&block)
-      assert_no_enqueued_jobs only: [ ActionMailer::DeliveryJob, ActionMailer::Parameterized::DeliveryJob ], &block
+      assert_enqueued_emails 0, &block
     end
+
+    private
+
+      def delivery_job_filter(job)
+        job_class = job.is_a?(Hash) ? job.fetch(:job) : job.class
+
+        Base.descendants.map(&:delivery_job).include?(job_class) ||
+        Base.descendants.map(&:parameterized_delivery_job).include?(job_class)
+      end
   end
 end
diff --git a/actionmailer/test/test_helper_test.rb b/actionmailer/test/test_helper_test.rb
index d31170706b..c700db9408 100644
--- a/actionmailer/test/test_helper_test.rb
+++ b/actionmailer/test/test_helper_test.rb
@@ -24,6 +24,17 @@ class TestHelperMailer < ActionMailer::Base
   end
 end
 
+class CustomDeliveryJob < ActionMailer::DeliveryJob
+end
+
+class CustomParameterizedDeliveryJob < ActionMailer::Parameterized::DeliveryJob
+end
+
+class CustomDeliveryMailer < TestHelperMailer
+  self.delivery_job = CustomDeliveryJob
+  self.parameterized_delivery_job = CustomParameterizedDeliveryJob
+end
+
 class TestHelperMailerTest < ActionMailer::TestCase
   include ActiveSupport::Testing::Stream
 
@@ -69,6 +80,26 @@ class TestHelperMailerTest < ActionMailer::TestCase
     end
   end
 
+  def test_assert_emails_with_custom_delivery_job
+    assert_nothing_raised do
+      assert_emails(1) do
+        silence_stream($stdout) do
+          CustomDeliveryMailer.test.deliver_later
+        end
+      end
+    end
+  end
+
+  def test_assert_emails_with_custom_parameterized_delivery_job
+    assert_nothing_raised do
+      assert_emails(1) do
+        silence_stream($stdout) do
+          CustomDeliveryMailer.with(foo: "bar").test_parameter_args.deliver_later
+        end
+      end
+    end
+  end
+
   def test_assert_emails_with_enqueued_emails
     assert_nothing_raised do
       assert_emails 1 do
@@ -201,6 +232,16 @@ class TestHelperMailerTest < ActionMailer::TestCase
     assert_match(/2 .* but 1/, error.message)
   end
 
+  def test_assert_enqueued_emails_with_custom_delivery_job
+    assert_nothing_raised do
+      assert_enqueued_emails(1) do
+        silence_stream($stdout) do
+          CustomDeliveryMailer.test.deliver_later
+        end
+      end
+    end
+  end
+
   def test_assert_enqueued_emails_too_many_sent
     error = assert_raise ActiveSupport::TestCase::Assertion do
       assert_enqueued_emails 1 do
@@ -252,6 +293,16 @@ class TestHelperMailerTest < ActionMailer::TestCase
     end
   end
 
+  def test_assert_enqueued_email_with_when_mailer_has_custom_delivery_job
+    assert_nothing_raised do
+      assert_enqueued_email_with CustomDeliveryMailer, :test do
+        silence_stream($stdout) do
+          CustomDeliveryMailer.test.deliver_later
+        end
+      end
+    end
+  end
+
   def test_assert_enqueued_email_with_with_no_block
     assert_nothing_raised do
       silence_stream($stdout) do
diff --git a/activejob/CHANGELOG.md b/activejob/CHANGELOG.md
index 8bbecd5a5a..e5cbff4a88 100644
--- a/activejob/CHANGELOG.md
+++ b/activejob/CHANGELOG.md
@@ -1,3 +1,8 @@
+*   Allow all assertion helpers that have a `only` and `except` keyword to accept
+    Procs.
+
+    *Edouard Chin*
+
 *   Restore HashWithIndifferentAccess support to ActiveJob::Arguments.deserialize.
 
     *Gannon McGibbon*
diff --git a/activejob/lib/active_job/queue_adapters/test_adapter.rb b/activejob/lib/active_job/queue_adapters/test_adapter.rb
index f73ad444ba..c134257ebc 100644
--- a/activejob/lib/active_job/queue_adapters/test_adapter.rb
+++ b/activejob/lib/active_job/queue_adapters/test_adapter.rb
@@ -65,11 +65,17 @@ module ActiveJob
 
         def filtered_job_class?(job)
           if filter
-            !Array(filter).include?(job.class)
+            !filter_as_proc(filter).call(job)
           elsif reject
-            Array(reject).include?(job.class)
+            filter_as_proc(reject).call(job)
           end
         end
+
+        def filter_as_proc(filter)
+          return filter if filter.is_a?(Proc)
+
+          ->(job) { Array(filter).include?(job.class) }
+        end
     end
   end
 end
diff --git a/activejob/lib/active_job/test_helper.rb b/activejob/lib/active_job/test_helper.rb
index 0deb68d0d2..e0a71f4da8 100644
--- a/activejob/lib/active_job/test_helper.rb
+++ b/activejob/lib/active_job/test_helper.rb
@@ -107,6 +107,9 @@ module ActiveJob
     #     end
     #   end
     #
+    # +:only+ and +:except+ options accepts Class, Array of Class or Proc. When passed a Proc,
+    # a hash containing the job's class and it's argument are passed as argument.
+    #
     # Asserts the number of times a job is enqueued to a specific queue by passing +:queue+ option.
     #
     #   def test_logging_job
@@ -163,6 +166,9 @@ module ActiveJob
     #     end
     #   end
     #
+    # +:only+ and +:except+ options accepts Class, Array of Class or Proc. When passed a Proc,
+    # a hash containing the job's class and it's argument are passed as argument.
+    #
     # Asserts that no jobs are enqueued to a specific queue by passing +:queue+ option
     #
     #   def test_no_logging
@@ -243,6 +249,18 @@ module ActiveJob
     #       end
     #     end
     #
+    # A proc may also be specified. When passed a Proc, the job's instance will be passed as argument.
+    #
+    #     def test_hello_and_logging_jobs
+    #       assert_nothing_raised do
+    #         assert_performed_jobs(1, only: ->(job) { job.is_a?(HelloJob) }) do
+    #           HelloJob.perform_later('jeremy')
+    #           LoggingJob.perform_later('stewie')
+    #           RescueJob.perform_later('david')
+    #         end
+    #       end
+    #     end
+    #
     # If the +:queue+ option is specified,
     # then only the job(s) enqueued to a specific queue will be performed.
     #
@@ -305,6 +323,9 @@ module ActiveJob
     #     end
     #   end
     #
+    # +:only+ and +:except+ options accepts Class, Array of Class or Proc. When passed a Proc,
+    # an instance of the job will be passed as argument.
+    #
     # If the +:queue+ option is specified,
     # then only the job(s) enqueued to a specific queue will not be performed.
     #
@@ -505,6 +526,9 @@ module ActiveJob
     #     assert_performed_jobs 1
     #   end
     #
+    # +:only+ and +:except+ options accepts Class, Array of Class or Proc. When passed a Proc,
+    # an instance of the job will be passed as argument.
+    #
     # If the +:queue+ option is specified,
     # then only the job(s) enqueued to a specific queue will be performed.
     #
@@ -569,9 +593,9 @@ module ActiveJob
           job_class = job.fetch(:job)
 
           if only
-            next false unless Array(only).include?(job_class)
+            next false unless filter_as_proc(only).call(job)
           elsif except
-            next false if Array(except).include?(job_class)
+            next false if filter_as_proc(except).call(job)
           end
 
           if queue
@@ -584,6 +608,12 @@ module ActiveJob
         end
       end
 
+      def filter_as_proc(filter)
+        return filter if filter.is_a?(Proc)
+
+        ->(job) { Array(filter).include?(job.fetch(:job)) }
+      end
+
       def enqueued_jobs_with(only: nil, except: nil, queue: nil, &block)
         jobs_with(enqueued_jobs, only: only, except: except, queue: queue, &block)
       end
diff --git a/activejob/test/cases/test_helper_test.rb b/activejob/test/cases/test_helper_test.rb
index 8dc32037ff..046033921d 100644
--- a/activejob/test/cases/test_helper_test.rb
+++ b/activejob/test/cases/test_helper_test.rb
@@ -114,6 +114,16 @@ class EnqueuedJobsTest < ActiveJob::TestCase
     end
   end
 
+  def test_assert_enqueued_jobs_with_only_option_as_proc
+    assert_nothing_raised do
+      assert_enqueued_jobs(1, only: ->(job) { job.fetch(:job).name == "HelloJob" }) do
+        HelloJob.perform_later("jeremy")
+        LoggingJob.perform_later
+        LoggingJob.perform_later
+      end
+    end
+  end
+
   def test_assert_enqueued_jobs_with_except_option
     assert_nothing_raised do
       assert_enqueued_jobs 1, except: LoggingJob do
@@ -124,6 +134,16 @@ class EnqueuedJobsTest < ActiveJob::TestCase
     end
   end
 
+  def test_assert_enqueued_jobs_with_except_option_as_proc
+    assert_nothing_raised do
+      assert_enqueued_jobs(1, except: ->(job) { job.fetch(:job).name == "LoggingJob" }) do
+        HelloJob.perform_later("jeremy")
+        LoggingJob.perform_later
+        LoggingJob.perform_later
+      end
+    end
+  end
+
   def test_assert_enqueued_jobs_with_only_and_except_option
     error = assert_raise ArgumentError do
       assert_enqueued_jobs 1, only: HelloJob, except: HelloJob do
@@ -911,6 +931,15 @@ class PerformedJobsTest < ActiveJob::TestCase
     end
   end
 
+  def test_assert_performed_jobs_with_only_option_as_proc
+    assert_nothing_raised do
+      assert_performed_jobs(1, only: ->(job) { job.is_a?(HelloJob) }) do
+        HelloJob.perform_later("jeremy")
+        LoggingJob.perform_later("bogdan")
+      end
+    end
+  end
+
   def test_assert_performed_jobs_without_block_with_only_option
     HelloJob.perform_later("jeremy")
     LoggingJob.perform_later("bogdan")
@@ -920,6 +949,15 @@ class PerformedJobsTest < ActiveJob::TestCase
     assert_performed_jobs 1, only: HelloJob
   end
 
+  def test_assert_performed_jobs_without_block_with_only_option_as_proc
+    HelloJob.perform_later("jeremy")
+    LoggingJob.perform_later("bogdan")
+
+    perform_enqueued_jobs
+
+    assert_performed_jobs(1, only: ->(job) { job.fetch(:job).name == "HelloJob" })
+  end
+
   def test_assert_performed_jobs_without_block_with_only_option_failure
     LoggingJob.perform_later("jeremy")
     LoggingJob.perform_later("bogdan")
@@ -942,6 +980,15 @@ class PerformedJobsTest < ActiveJob::TestCase
     end
   end
 
+  def test_assert_performed_jobs_with_except_option_as_proc
+    assert_nothing_raised do
+      assert_performed_jobs(1, except: ->(job) { job.is_a?(HelloJob) }) do
+        HelloJob.perform_later("jeremy")
+        LoggingJob.perform_later("bogdan")
+      end
+    end
+  end
+
   def test_assert_performed_jobs_without_block_with_except_option
     HelloJob.perform_later("jeremy")
     LoggingJob.perform_later("bogdan")
@@ -951,6 +998,15 @@ class PerformedJobsTest < ActiveJob::TestCase
     assert_performed_jobs 1, except: HelloJob
   end
 
+  def test_assert_performed_jobs_without_block_with_except_option_as_proc
+    HelloJob.perform_later("jeremy")
+    LoggingJob.perform_later("bogdan")
+
+    perform_enqueued_jobs
+
+    assert_performed_jobs(1, except: ->(job) { job.fetch(:job).name == "HelloJob" })
+  end
+
   def test_assert_performed_jobs_without_block_with_except_option_failure
     HelloJob.perform_later("jeremy")
     HelloJob.perform_later("bogdan")