Fix ActionMailer assertion not working for mail defining delivery_job:

- If a Mail defines a custom delivery_job, all ActionMailer assertion
  helper (assert_emails, assert_enqueued_emails ...) wouldn't work.

  ```ruby
    MyMailer < ApplicationMailer
      self.delivery_job = MyJob
    end

    # This assertion will fail
    assert_emails(1) do
      MyMailer.my_mail.deliver_later
    end

  This PR leverage the new ActiveJob feature that accepts Procs for the
  `only` keyword and check if the delivery job is one of ActionMailer
   registered ones.

3 files changed, 70 insertions, 5 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