Merge pull request #754 from Shopify/add-after-environment-load-hook

Add Spring.after_environment_load hook for post-Rails preloading

Author: rafaelfranca

lib/spring/application.rb

@@ -116,6 +116,8 @@ def preload
 
       require Spring.application_root_path.join("config", "environment")
 
+      invoke_after_environment_load_callbacks
+
       disconnect_database
 
       @preloaded = :success
@@ -300,6 +302,12 @@ def invoke_after_fork_callbacks
       end
     end
 
+    def invoke_after_environment_load_callbacks
+      Spring.after_environment_load_callbacks.each do |callback|
+        callback.call
+      end
+    end
+
     def loaded_application_features
       root = Spring.application_root_path.to_s
       $LOADED_FEATURES.select { |f| f.start_with?(root) }

lib/spring/configuration.rb

@@ -35,6 +35,14 @@ def after_fork(&block)
       after_fork_callbacks << block
     end
 
+    def after_environment_load_callbacks
+      @after_environment_load_callbacks ||= []
+    end
+
+    def after_environment_load(&block)
+      after_environment_load_callbacks << block
+    end
+
     def spawn_on_env
       @spawn_on_env ||= []
     end

test/support/acceptance_test.rb

@@ -218,6 +218,22 @@ def self.omg
         assert_success app.spring_test_command
       end
 
+      test "app gets reloaded when after_environment_load preloads test_helper" do
+        File.write(app.spring_config, <<-RUBY.strip_heredoc)
+          Spring.after_environment_load do
+            require File.expand_path("../test/test_helper", __dir__)
+          end
+        RUBY
+
+        assert_success app.spring_test_command
+
+        test_helper = app.path("test/test_helper.rb")
+        File.write(test_helper, test_helper.read + "\nraise 'omg'\n")
+
+        app.await_reload
+        assert_failure app.spring_test_command, stderr: /omg \(RuntimeError\)/, log: /child \d+ shutdown/
+      end
+
       test "app recovers when a boot-level error is introduced" do
         config = app.application_config.read
 
@@ -531,6 +547,13 @@ def exec_name
         assert_success "bin/rails runner 'puts 2'", stdout: "!callback!\n2"
       end
 
+      test "after environment load callback" do
+        File.write(app.spring_config, "Spring.after_environment_load { puts '!callback!' }")
+
+        assert_success app.spring_test_command, stdout: "!callback!"
+        refute_output_includes app.spring_test_command, stdout: "!callback!"
+      end
+
       test "global config file evaluated" do
         File.write("#{app.user_home}/.spring.rb", "Spring.after_fork { puts '!callback!' }")
         assert_success "bin/rails runner 'puts 2'", stdout: "!callback!\n2"

test/unit/configuration_test.rb

@@ -0,0 +1,58 @@
+require_relative "../helper"
+require "spring/configuration"
+
+class ConfigurationTest < ActiveSupport::TestCase
+  test "after_environment_load_callbacks is empty by default" do
+    assert_equal [], Spring.after_environment_load_callbacks
+  end
+
+  test "after_environment_load registers a callback" do
+    callbacks = []
+    Spring.stub(:after_environment_load_callbacks, callbacks) do
+      Spring.after_environment_load { true }
+    end
+    assert_equal 1, callbacks.size
+  end
+
+  test "after_environment_load accumulates multiple callbacks in order" do
+    callbacks = []
+    Spring.stub(:after_environment_load_callbacks, callbacks) do
+      Spring.after_environment_load { :first }
+      Spring.after_environment_load { :second }
+    end
+    assert_equal [:first, :second], callbacks.map(&:call)
+  end
+
+  test "after_environment_load callbacks are independent from after_fork callbacks" do
+    env_callbacks  = []
+    fork_callbacks = []
+    Spring.stub(:after_environment_load_callbacks, env_callbacks) do
+      Spring.stub(:after_fork_callbacks, fork_callbacks) do
+        Spring.after_environment_load { true }
+      end
+    end
+    assert_equal 1, env_callbacks.size
+    assert_equal 0, fork_callbacks.size
+  end
+
+  test "after_fork_callbacks is empty by default" do
+    assert_equal [], Spring.after_fork_callbacks
+  end
+
+  test "after_fork registers a callback" do
+    callbacks = []
+    Spring.stub(:after_fork_callbacks, callbacks) do
+      Spring.after_fork { true }
+    end
+    assert_equal 1, callbacks.size
+  end
+
+  test "after_fork accumulates multiple callbacks in order" do
+    callbacks = []
+    Spring.stub(:after_fork_callbacks, callbacks) do
+      Spring.after_fork { :first }
+      Spring.after_fork { :second }
+    end
+    assert_equal [:first, :second], callbacks.map(&:call)
+  end
+end