Merge pull request #51525 from Shopify/aj-after-commit-return-value

Fix `ActiveJob::EnqueueAfterTransactionCommit` API

activejob/lib/active_job/enqueue_after_transaction_commit.rb

@@ -16,9 +16,11 @@ module EnqueueAfterTransactionCommit # :nodoc:
       #  - `:never` forces the job to be queued immediately.
       #  - `:default` lets the queue adapter define the behavior (recommended).
       class_attribute :enqueue_after_transaction_commit, instance_accessor: false, instance_predicate: false, default: :never
+    end
 
-      around_enqueue do |job, block|
-        after_transaction = case job.class.enqueue_after_transaction_commit
+    private
+      def raw_enqueue
+        after_transaction = case self.class.enqueue_after_transaction_commit
         when :always
           true
         when :never
@@ -28,10 +30,14 @@ module EnqueueAfterTransactionCommit # :nodoc:
         end
 
         if after_transaction
-          ActiveRecord.after_all_transactions_commit(&block)
-        else
-          block.call
+          self.successfully_enqueued = true
+          ActiveRecord.after_all_transactions_commit do
+            self.successfully_enqueued = false
+            super
           end
+          self
+        else
+          super
         end
       end
   end

activejob/lib/active_job/enqueuing.rb

@@ -50,14 +50,17 @@ module ClassMethods
       # custom serializers.
       #
       # Returns an instance of the job class queued with arguments available in
-      # Job#arguments or false if the enqueue did not succeed.
+      # Job#arguments or +false+ if the enqueue did not succeed.
       #
       # After the attempted enqueue, the job will be yielded to an optional block.
       #
       # If Active Job is used conjointly with Active Record, and #perform_later is called
       # inside an Active Record transaction, then the enqueue is implicitly deferred to after
-      # the transaction is committed, or droped if it's rolled back. This behavior can
-      # be changed on a per job basis:
+      # the transaction is committed, or droped if it's rolled back. In such case #perform_later
+      # will return the job instance like if it was successfully enqueued, but will still return
+      # +false+ if a callback prevented the job from being enqueued.
+      #
+      # This behavior can be changed on a per job basis:
       #
       #  class NotificationJob < ApplicationJob
       #    self.enqueue_after_transaction_commit = false
@@ -98,6 +101,18 @@ def enqueue(options = {})
       self.successfully_enqueued = false
 
       run_callbacks :enqueue do
+        raw_enqueue
+      end
+
+      if successfully_enqueued?
+        self
+      else
+        false
+      end
+    end
+
+    private
+      def raw_enqueue
         if scheduled_at
           queue_adapter.enqueue_at self, scheduled_at.to_f
         else
@@ -108,12 +123,5 @@ def enqueue(options = {})
       rescue EnqueueError => e
         self.enqueue_error = e
       end
-
-      if successfully_enqueued?
-        self
-      else
-        false
-      end
-    end
   end
 end

activejob/test/cases/enqueue_after_transaction_commit_test.rb

@@ -0,0 +1,105 @@
+# frozen_string_literal: true
+
+require "helper"
+require "jobs/enqueue_error_job"
+
+class EnqueueAfterTransactionCommitTest < ActiveSupport::TestCase
+  class FakeActiveRecord
+    attr_reader :calls
+
+    def initialize(should_yield = true)
+      @calls = 0
+      @yield = should_yield
+      @callbacks = []
+    end
+
+    def after_all_transactions_commit(&block)
+      @calls += 1
+      if @yield
+        yield
+      else
+        @callbacks << block
+      end
+    end
+
+    def run_after_commit_callbacks
+      callbacks, @callbacks = @callbacks, []
+      callbacks.each(&:call)
+    end
+  end
+
+  class EnqueueAfterCommitJob < ActiveJob::Base
+    self.enqueue_after_transaction_commit = :always
+
+    def perform
+      # noop
+    end
+  end
+
+  class ErrorEnqueueAfterCommitJob < EnqueueErrorJob
+    class EnqueueErrorAdapter
+      def enqueue_after_transaction_commit?
+        true
+      end
+
+      def enqueue(...)
+        raise ActiveJob::EnqueueError, "There was an error enqueuing the job"
+      end
+
+      def enqueue_at(...)
+        raise ActiveJob::EnqueueError, "There was an error enqueuing the job"
+      end
+    end
+
+    self.queue_adapter = EnqueueErrorAdapter.new
+    self.enqueue_after_transaction_commit = :always
+
+    def perform
+      # noop
+    end
+  end
+
+  test "#perform_later wait for transactions to complete before enqueuing the job" do
+    fake_active_record = FakeActiveRecord.new
+    stub_const(Object, :ActiveRecord, fake_active_record, exists: false) do
+      assert_difference -> { fake_active_record.calls }, +1 do
+        EnqueueAfterCommitJob.perform_later
+      end
+    end
+  end
+
+  test "#perform_later returns the Job instance even if it's delayed by `after_all_transactions_commit`" do
+    fake_active_record = FakeActiveRecord.new(false)
+    stub_const(Object, :ActiveRecord, fake_active_record, exists: false) do
+      job = EnqueueAfterCommitJob.perform_later
+      assert_instance_of EnqueueAfterCommitJob, job
+      assert_predicate job, :successfully_enqueued?
+    end
+  end
+
+  test "#perform_later yields the enqueued Job instance even if it's delayed by `after_all_transactions_commit`" do
+    fake_active_record = FakeActiveRecord.new(false)
+    stub_const(Object, :ActiveRecord, fake_active_record, exists: false) do
+      called = false
+      job = EnqueueAfterCommitJob.perform_later do |yielded_job|
+        called = true
+        assert_instance_of EnqueueAfterCommitJob, yielded_job
+      end
+      assert called, "#perform_later yielded the job"
+      assert_instance_of EnqueueAfterCommitJob, job
+      assert_predicate job, :successfully_enqueued?
+    end
+  end
+
+  test "#perform_later assumes successful enqueue, but update status later" do
+    fake_active_record = FakeActiveRecord.new(false)
+    stub_const(Object, :ActiveRecord, fake_active_record, exists: false) do
+      job = ErrorEnqueueAfterCommitJob.perform_later
+      assert_instance_of ErrorEnqueueAfterCommitJob, job
+      assert_predicate job, :successfully_enqueued?
+
+      fake_active_record.run_after_commit_callbacks
+      assert_not_predicate job, :successfully_enqueued?
+    end
+  end
+end

activejob/test/helper.rb

@@ -23,3 +23,5 @@ def adapter_is?(*adapter_class_symbols)
 end
 
 require_relative "../../tools/test_common"
+
+ActiveJob::Base.include(ActiveJob::EnqueueAfterTransactionCommit)

activesupport/CHANGELOG.md

@@ -1,3 +1,7 @@
+*   `stub_const` now accepts a `exists: false` parameter to allow stubbing missing constants.
+
+    *Jean Boussier*
+
 *   Make ActiveSupport::BacktraceCleaner copy filters and silencers on dup and clone
 
     Previously the copy would still share the internal silencers and filters array,

activesupport/lib/active_support/testing/constant_stubbing.rb

@@ -15,10 +15,19 @@ module ConstantStubbing
       # Using this method rather than forcing <tt>World::List::Import::LARGE_IMPORT_THRESHOLD = 5000</tt> prevents
       # warnings from being thrown, and ensures that the old value is returned after the test has completed.
       #
+      # If the constant doesn't already exists, but you need it set for the duration of the block
+      # you can do so by passing `exists: false`.
+      #
+      #   stub_const(object, :SOME_CONST, 1, exists: false) do
+      #     assert_equal 1, SOME_CONST
+      #   end
+      #
       # Note: Stubbing a const will stub it across all threads. So if you have concurrent threads
       # (like separate test suites running in parallel) that all depend on the same constant, it's possible
       # divergent stubbing will trample on each other.
-      def stub_const(mod, constant, new_value)
+      def stub_const(mod, constant, new_value, exists: true)
+        if exists
+          begin
             old_value = mod.const_get(constant, false)
             mod.send(:remove_const, constant)
             mod.const_set(constant, new_value)
@@ -27,6 +36,19 @@ def stub_const(mod, constant, new_value)
             mod.send(:remove_const, constant)
             mod.const_set(constant, old_value)
           end
+        else
+          if mod.const_defined?(constant)
+            raise NameError, "already defined constant #{constant} in #{mod.name}"
+          end
+
+          begin
+            mod.const_set(constant, new_value)
+            yield
+          ensure
+            mod.send(:remove_const, constant)
+          end
+        end
+      end
     end
   end
 end

activesupport/test/test_case_test.rb

@@ -580,7 +580,7 @@ class TestConstStubbing < ActiveSupport::TestCase
     assert_equal 1, ConstStubbable::CONSTANT
   end
 
-  test "trying to stub a constant that does not exist in the receiver raises NameError" do
+  test "stubbing a constant that does not exist in the receiver raises NameError" do
     assert_raises(NameError) do
       stub_const(ConstStubbable, :NOT_A_CONSTANT, 1) { }
     end
@@ -589,4 +589,18 @@ class TestConstStubbing < ActiveSupport::TestCase
       stub_const(SubclassOfConstStubbable, :CONSTANT, 1) { }
     end
   end
+
+  test "stubbing a constant that does not exist can be done with `exists: false`" do
+    stub_const(ConstStubbable, :NOT_A_CONSTANT, 1, exists: false) do
+      assert_equal 1, ConstStubbable::NOT_A_CONSTANT
+    end
+
+    assert_raises(NameError) do
+      ConstStubbable::NOT_A_CONSTANT
+    end
+
+    assert_raises(NameError) do
+      stub_const(Object, :ConstStubbable, 1, exists: false)
+    end
+  end
 end