From 3606811821820dda794938a263a0cdc87385fc76 Mon Sep 17 00:00:00 2001
From: Matthew Erhard <matthew.erhard@egov.com>
Date: Mon, 14 Jul 2014 14:42:54 -0400
Subject: [PATCH] allow arbitrary binstub prelude code to be inserted into a
 binstub

---
 lib/spring/client/binstub.rb  |  8 +++++++-
 lib/spring/command_wrapper.rb |  6 ++++++
 test/acceptance/app_test.rb   | 19 +++++++++++++++++++
 3 files changed, 32 insertions(+), 1 deletion(-)

diff --git a/lib/spring/client/binstub.rb b/lib/spring/client/binstub.rb
index dd3628ff..9bad9f30 100644
--- a/lib/spring/client/binstub.rb
+++ b/lib/spring/client/binstub.rb
@@ -104,8 +104,14 @@ def generate(fallback = nil)
             fallback = "require 'bundler/setup'\n" \
                        "load Gem.bin_path('#{command.gem_name}', '#{command.exec_name}')\n"
           end
+          if prelude = command.binstub_prelude
+            formatted_prelude = prelude.chomp.gsub(/^(?!$)/, '  ')
+            loader = LOADER.sub(/^end$/, "else\n#{formatted_prelude}\nend")
+          else
+            loader = LOADER
+          end
 
-          File.write(command.binstub, "#!/usr/bin/env ruby\n#{LOADER}#{fallback}")
+          File.write(command.binstub, "#!/usr/bin/env ruby\n#{loader}#{fallback}")
           command.binstub.chmod 0755
         end
 
diff --git a/lib/spring/command_wrapper.rb b/lib/spring/command_wrapper.rb
index 539a0778..98402bb2 100644
--- a/lib/spring/command_wrapper.rb
+++ b/lib/spring/command_wrapper.rb
@@ -69,6 +69,12 @@ def binstub_name
       "bin/#{name}"
     end
 
+    def binstub_prelude
+      if command.respond_to?(:binstub_prelude)
+        command.binstub_prelude
+      end
+    end
+
     def exec
       if binstub.exist?
         binstub.to_s
diff --git a/test/acceptance/app_test.rb b/test/acceptance/app_test.rb
index c1633f98..37ad5960 100644
--- a/test/acceptance/app_test.rb
+++ b/test/acceptance/app_test.rb
@@ -188,6 +188,25 @@ def exec_name
     assert_success "bin/rake -T", stdout: "rake db:migrate"
   end
 
+  test "binstub with prelude code" do
+    prelude = "Prelude code line 1\nPrelude code line 2\n"
+
+    File.write(app.spring_config, <<-CODE)
+      class PreludeCode
+        def binstub_prelude
+          "#{prelude}"
+        end
+      end
+
+      Spring.register_command "prelude", PreludeCode.new
+    CODE
+
+    assert_success "bin/spring binstub prelude"
+    prelude.each_line do |line|
+      assert app.path("bin/prelude").read.include?(line), "'#{line}' not found in bin/prelude"
+    end
+  end
+
   test "binstub when spring is uninstalled" do
     app.run! "gem uninstall --ignore-dependencies spring"
     File.write(app.gemfile, app.gemfile.read.gsub(/gem 'spring.*/, ""))