Fix upgrade/remove with new binstub variations

lib/spring/client/binstub.rb

@@ -48,6 +48,12 @@ class Binstub < Command
 
       OLD_BINSTUB = %{if !Process.respond_to?(:fork) || Gem::Specification.find_all_by_name("spring").empty?}
 
+      BINSTUB_VARIATIONS = Regexp.union [
+        %{begin\n  load File.expand_path("../spring", __FILE__)\nrescue LoadError\nend\n},
+        %{begin\n  load File.expand_path('../spring', __FILE__)\nrescue LoadError\nend\n},
+        LOADER
+      ]
+
       class Item
         attr_reader :command, :existing
 
@@ -74,8 +80,12 @@ def add
               fallback = nil if fallback.include?("exec")
               generate(fallback)
               status "upgraded"
-            elsif existing =~ /load .*spring/
+            elsif existing.include?(LOADER)
               status "spring already present"
+            elsif existing =~ BINSTUB_VARIATIONS
+              upgraded = existing.sub(BINSTUB_VARIATIONS, LOADER)
+              File.write(command.binstub, upgraded)
+              status "upgraded"
             else
               head, shebang, tail = existing.partition(SHEBANG)
 
@@ -110,7 +120,7 @@ def generate(fallback = nil)
 
         def remove
           if existing
-            File.write(command.binstub, existing.sub(LOADER, ""))
+            File.write(command.binstub, existing.sub(BINSTUB_VARIATIONS, ""))
             status "spring removed"
           end
         end
@@ -119,7 +129,7 @@ def remove
       attr_reader :bindir, :items
 
       def self.description
-        "Generate spring based binstubs. Use --all to generate a binstub for all known commands."
+        "Generate spring based binstubs. Use --all to generate a binstub for all known commands. Use --remove to revert."
       end
 
       def self.rails_command

lib/spring/test/acceptance_test.rb

@@ -208,6 +208,11 @@ def exec_name
         assert_success "bin/rake -T", stdout: "rake db:migrate"
       end
 
+      test "binstub remove all" do
+        assert_success "bin/spring binstub --remove --all"
+        refute File.exist?(app.path("bin/spring"))
+      end
+
       test "binstub when spring gem is missing" do
         without_gem "spring-#{Spring::VERSION}" do
           File.write(app.gemfile, app.gemfile.read.gsub(/gem 'spring.*/, ""))
@@ -224,7 +229,7 @@ def exec_name
         end
       end
 
-      test "binstub upgrade" do
+      test "binstub upgrade with old binstub" do
         File.write(app.path("bin/rake"), <<-RUBY.strip_heredoc)
           #!/usr/bin/env ruby
 
@@ -269,6 +274,93 @@ def exec_name
         assert_equal expected, app.path("bin/rails").read
       end
 
+      test "binstub upgrade with new binstub variations" do
+        # older variation with double quotes
+        File.write(app.path("bin/rake"), <<-RUBY.strip_heredoc)
+          #!/usr/bin/env ruby
+          begin
+            load File.expand_path("../spring", __FILE__)
+          rescue LoadError
+          end
+          require 'bundler/setup'
+          load Gem.bin_path('rake', 'rake')
+        RUBY
+
+        # newer variation with single quotes
+        File.write(app.path("bin/rails"), <<-RUBY.strip_heredoc)
+          #!/usr/bin/env ruby
+          begin
+            load File.expand_path('../spring', __FILE__)
+          rescue LoadError
+          end
+          APP_PATH = File.expand_path('../../config/application',  __FILE__)
+          require_relative '../config/boot'
+          require 'rails/commands'
+        RUBY
+
+        assert_success "bin/spring binstub --all", stdout: "upgraded"
+
+        expected = <<-RUBY.gsub(/^          /, "")
+          #!/usr/bin/env ruby
+          #{Spring::Client::Binstub::LOADER.strip}
+          require 'bundler/setup'
+          load Gem.bin_path('rake', 'rake')
+        RUBY
+        assert_equal expected, app.path("bin/rake").read
+
+        expected = <<-RUBY.gsub(/^          /, "")
+          #!/usr/bin/env ruby
+          #{Spring::Client::Binstub::LOADER.strip}
+          APP_PATH = File.expand_path('../../config/application',  __FILE__)
+          require_relative '../config/boot'
+          require 'rails/commands'
+        RUBY
+        assert_equal expected, app.path("bin/rails").read
+      end
+
+      test "binstub remove with new binstub variations" do
+        # older variation with double quotes
+        File.write(app.path("bin/rake"), <<-RUBY.strip_heredoc)
+          #!/usr/bin/env ruby
+          begin
+            load File.expand_path("../spring", __FILE__)
+          rescue LoadError
+          end
+          require 'bundler/setup'
+          load Gem.bin_path('rake', 'rake')
+        RUBY
+
+        # newer variation with single quotes
+        File.write(app.path("bin/rails"), <<-RUBY.strip_heredoc)
+          #!/usr/bin/env ruby
+          begin
+            load File.expand_path('../spring', __FILE__)
+          rescue LoadError
+          end
+          APP_PATH = File.expand_path('../../config/application',  __FILE__)
+          require_relative '../config/boot'
+          require 'rails/commands'
+        RUBY
+
+        assert_success "bin/spring binstub --remove rake", stdout: "bin/rake: spring removed"
+        assert_success "bin/spring binstub --remove rails", stdout: "bin/rails: spring removed"
+
+        expected = <<-RUBY.strip_heredoc
+          #!/usr/bin/env ruby
+          require 'bundler/setup'
+          load Gem.bin_path('rake', 'rake')
+        RUBY
+        assert_equal expected, app.path("bin/rake").read
+
+        expected = <<-RUBY.strip_heredoc
+          #!/usr/bin/env ruby
+          APP_PATH = File.expand_path('../../config/application',  __FILE__)
+          require_relative '../config/boot'
+          require 'rails/commands'
+        RUBY
+        assert_equal expected, app.path("bin/rails").read
+      end
+
       test "after fork callback" do
         File.write(app.spring_config, "Spring.after_fork { puts '!callback!' }")
         assert_success "bin/rails runner 'puts 2'", stdout: "!callback!\n2"