Merge pull request rails#55247 from byroot/nested-state

Allow for nested ExecutionContext in test

Author: byroot

activestorage/test/test_helper.rb

@@ -12,6 +12,7 @@
 require "active_support/testing/autorun"
 require "image_processing/mini_magick"
 
+require "active_support/current_attributes/test_helper"
 require "active_record/testing/query_assertions"
 
 require "active_job"
@@ -25,6 +26,7 @@
 class ActiveSupport::TestCase
   self.file_fixture_path = ActiveStorage::FixtureSet.file_fixture_path
 
+  include ActiveSupport::CurrentAttributes::TestHelper
   include ActiveRecord::TestFixtures
   include ActiveRecord::Assertions::QueryAssertions
 
@@ -34,10 +36,6 @@ class ActiveSupport::TestCase
     ActiveStorage::Current.url_options = { protocol: "https://", host: "example.com", port: nil }
   end
 
-  teardown do
-    ActiveStorage::Current.reset
-  end
-
   private
     def create_blob(key: nil, data: "Hello world!", filename: "hello.txt", content_type: "text/plain", identify: true, service_name: nil, record: nil)
       ActiveStorage::Blob.create_and_upload! key: key, io: StringIO.new(data), filename: filename, content_type: content_type, identify: identify, service_name: service_name, record: record

activesupport/CHANGELOG.md

@@ -1,3 +1,21 @@
+*   Improve `CurrentAttribute` and `ExecutionContext` state managment in test cases.
+
+    Previously these two global state would be entirely cleared out whenever calling
+    into code that is wrapped by the Rails executor, typically Action Controller or
+    Active Job helpers:
+
+    ```ruby
+    test "#index works" do
+      CurrentUser.id = 42
+      get :index
+      CurrentUser.id == nil
+    end
+    ```
+
+    Now re-entering the executor properly save and restore that state.
+
+    *Jean Boussier*
+
 *   The new method `ActiveSupport::BacktraceCleaner#first_clean_location`
     returns the first clean location of the caller's call stack, or `nil`.
     Locations are `Thread::Backtrace::Location` objects. Useful when you want to

activesupport/lib/active_support/current_attributes.rb

@@ -1,6 +1,7 @@
 # frozen_string_literal: true
 
 require "active_support/callbacks"
+require "active_support/execution_context"
 require "active_support/core_ext/object/with"
 require "active_support/core_ext/enumerable"
 require "active_support/core_ext/module/delegation"
@@ -154,8 +155,10 @@ def resets(*methods, &block)
       delegate :set, :reset, to: :instance
 
       def clear_all # :nodoc:
-        current_instances.each_value(&:reset)
-        current_instances.clear
+        if instances = current_instances
+          instances.values.each(&:reset)
+          instances.clear
+        end
       end
 
       private
@@ -164,7 +167,7 @@ def generated_attribute_methods
         end
 
         def current_instances
-          IsolatedExecutionState[:current_attributes_instances] ||= {}
+          ExecutionContext.current_attributes_instances
         end
 
         def current_instances_key

activesupport/lib/active_support/execution_context.rb

@@ -2,8 +2,41 @@
 
 module ActiveSupport
   module ExecutionContext # :nodoc:
+    class Record
+      attr_reader :store, :current_attributes_instances
+
+      def initialize
+        @store = {}
+        @current_attributes_instances = {}
+        @stack = []
+      end
+
+      def push
+        @stack << @store << @current_attributes_instances
+        @store = {}
+        @current_attributes_instances = {}
+        self
+      end
+
+      def pop
+        @current_attributes_instances = @stack.pop
+        @store = @stack.pop
+        self
+      end
+    end
+
     @after_change_callbacks = []
+
+    # Execution context nesting should only legitimately happen during test
+    # because the test case itself is wrapped in an executor, and it might call
+    # into a controller or job which should be executed with their own fresh context.
+    # However in production this should never happen, and for extra safety we make sure to
+    # fully clear the state at the end of the request or job cycle.
+    @nestable = false
+
     class << self
+      attr_accessor :nestable
+
       def after_change(&block)
         @after_change_callbacks << block
       end
@@ -14,9 +47,11 @@ def set(**options)
         options.symbolize_keys!
         keys = options.keys
 
-        store = self.store
+        store = record.store
 
-        previous_context = keys.zip(store.values_at(*keys)).to_h
+        previous_context = if block_given?
+          keys.zip(store.values_at(*keys)).to_h
+        end
 
         store.merge!(options)
         @after_change_callbacks.each(&:call)
@@ -32,21 +67,43 @@ def set(**options)
       end
 
       def []=(key, value)
-        store[key.to_sym] = value
+        record.store[key.to_sym] = value
         @after_change_callbacks.each(&:call)
       end
 
       def to_h
-        store.dup
+        record.store.dup
+      end
+
+      def push
+        if @nestable
+          record.push
+        else
+          clear
+        end
+        self
+      end
+
+      def pop
+        if @nestable
+          record.pop
+        else
+          clear
+        end
+        self
       end
 
       def clear
-        store.clear
+        IsolatedExecutionState[:active_support_execution_context] = nil
+      end
+
+      def current_attributes_instances
+        record.current_attributes_instances
       end
 
       private
-        def store
-          IsolatedExecutionState[:active_support_execution_context] ||= {}
+        def record
+          IsolatedExecutionState[:active_support_execution_context] ||= Record.new
         end
     end
   end

activesupport/lib/active_support/railtie.rb

@@ -39,18 +39,24 @@ class Railtie < Rails::Railtie # :nodoc:
     end
 
     initializer "active_support.reset_execution_context" do |app|
-      app.reloader.before_class_unload { ActiveSupport::ExecutionContext.clear }
-      app.executor.to_run              { ActiveSupport::ExecutionContext.clear }
-      app.executor.to_complete         { ActiveSupport::ExecutionContext.clear }
-    end
+      app.reloader.before_class_unload do
+        ActiveSupport::CurrentAttributes.clear_all
+        ActiveSupport::ExecutionContext.clear
+      end
+
+      app.executor.to_run do
+        ActiveSupport::ExecutionContext.push
+      end
 
-    initializer "active_support.clear_all_current_attributes_instances" do |app|
-      app.reloader.before_class_unload { ActiveSupport::CurrentAttributes.clear_all }
-      app.executor.to_run              { ActiveSupport::CurrentAttributes.clear_all }
-      app.executor.to_complete         { ActiveSupport::CurrentAttributes.clear_all }
+      app.executor.to_complete do
+        ActiveSupport::CurrentAttributes.clear_all
+        ActiveSupport::ExecutionContext.pop
+      end
 
       ActiveSupport.on_load(:active_support_test_case) do
         if app.config.active_support.executor_around_test_case
+          ActiveSupport::ExecutionContext.nestable = true
+
           require "active_support/executor/test_helper"
           include ActiveSupport::Executor::TestHelper
         else

activesupport/test/current_attributes_test.rb

@@ -66,9 +66,6 @@ class Session < ActiveSupport::CurrentAttributes
     attribute :current, :previous
   end
 
-  # Eagerly set-up `instance`s by reference.
-  [ Current.instance, Session.instance ]
-
   # Use library specific minitest hook to catch Time.zone before reset is called via TestHelper
   def before_setup
     @original_time_zone = Time.zone
@@ -307,4 +304,50 @@ def self.name
     assert_equal [], current.method(:attr).parameters
     assert_equal [[:req, :value]], current.method(:attr=).parameters
   end
+
+
+  test "set and restore attributes when re-entering the executor" do
+    ActiveSupport::ExecutionContext.with(nestable: true) do
+      # simulate executor hooks from active_support/railtie.rb
+      executor = Class.new(ActiveSupport::Executor)
+      executor.to_run do
+        ActiveSupport::ExecutionContext.push
+      end
+
+      executor.to_complete do
+        ActiveSupport::CurrentAttributes.clear_all
+        ActiveSupport::ExecutionContext.pop
+      end
+
+      Current.world = "world/1"
+      Current.account = "account/1"
+
+      assert_equal "world/1", Current.world
+      assert_equal "account/1", Current.account
+
+      Current.set(world: "world/2", account: "account/2") do
+        assert_equal "world/2", Current.world
+        assert_equal "account/2", Current.account
+
+        executor.wrap do
+          assert_nil Current.world
+          assert_nil Current.account
+
+          Current.world = "world/3"
+          Current.account = "account/3"
+
+          assert_equal "world/3", Current.world
+          assert_equal "account/3", Current.account
+
+          ActiveSupport::CurrentAttributes.clear_all
+
+          assert_nil Current.world
+          assert_nil Current.account
+        end
+      end
+
+      assert_equal "world/1", Current.world
+      assert_equal "account/1", Current.account
+    end
+  end
 end