Refactor tests that switch RAILS_ENV and RACK_ENV

This cleanup aims to fix a build failure:
https://travis-ci.org/rails/rails/jobs/3515951/#L482

Since travis always have both ENV vars set to "test", a test is failing
where it's expected to output the default env "development", but "test"
is the result due to RACK_ENV being set when we expect it to not be.

By cleaning this duplication we ensure that changing any of these env
variables will pick the right expected value.

railties/test/abstract_unit.rb

@@ -7,7 +7,6 @@
 require 'fileutils'
 
 require 'active_support'
-
 require 'action_controller'
 require 'rails/all'
 

railties/test/application/runner_test.rb

@@ -1,8 +1,10 @@
 require 'isolation/abstract_unit'
+require 'env_helpers'
 
 module ApplicationTests
   class RunnerTest < ActiveSupport::TestCase
     include ActiveSupport::Testing::Isolation
+    include EnvHelpers
 
     def setup
       build_app
@@ -73,21 +75,15 @@ def test_default_environment
     end
 
     def test_environment_with_rails_env
-      orig = ENV['RAILS_ENV']
-      ENV['RAILS_ENV'] = "production"
-      assert_match "production", Dir.chdir(app_path) { `bundle exec rails runner "puts Rails.env"` }
-    ensure
-      ENV['RAILS_ENV'] = orig
+      with_rails_env "production" do
+        assert_match "production", Dir.chdir(app_path) { `bundle exec rails runner "puts Rails.env"` }
+      end
     end
 
     def test_environment_with_rack_env
-      rack, rails = ENV['RACK_ENV'], ENV['RAILS_ENV']
-      ENV['RACK_ENV'] = "production"
-      ENV['RAILS_ENV'] = nil
-      assert_match "production", Dir.chdir(app_path) { `bundle exec rails runner "puts Rails.env"` }
-    ensure
-      ENV['RAILS_ENV'] = rails
-      ENV['RACK_ENV'] = rack
+      with_rack_env "production" do
+        assert_match "production", Dir.chdir(app_path) { `bundle exec rails runner "puts Rails.env"` }
+      end
     end
   end
 end

railties/test/commands/console_test.rb

@@ -1,14 +1,14 @@
 require 'abstract_unit'
+require 'env_helpers'
 require 'rails/commands/console'
 
 class Rails::ConsoleTest < ActiveSupport::TestCase
+  include EnvHelpers
+
   class FakeConsole
     def self.start; end
   end
 
-  def setup
-  end
-
   def test_sandbox_option
     console = Rails::Console.new(app, parse_arguments(["--sandbox"]))
     assert console.sandbox?
@@ -85,7 +85,7 @@ def test_default_environment_with_rack_env
       assert_match(/\sproduction\s/, output)
     end
   end
-  
+
   def test_e_option
     start ['-e', 'special-production']
     assert_match(/\sspecial-production\s/, output)
@@ -133,20 +133,4 @@ def app
   def parse_arguments(args)
     Rails::Console.parse_arguments(args)
   end
-
-  def with_rails_env(env)
-    rails = ENV['RAILS_ENV']
-    ENV['RAILS_ENV'] = env
-    yield
-  ensure
-    ENV['RAILS_ENV'] = rails
-  end
-
-  def with_rack_env(env)
-    rack = ENV['RACK_ENV']
-    ENV['RACK_ENV'] = env
-    with_rails_env(nil) { yield }
-  ensure
-    ENV['RACK_ENV'] = rack
-  end
 end

railties/test/commands/server_test.rb

@@ -1,7 +1,9 @@
 require 'abstract_unit'
+require 'env_helpers'
 require 'rails/commands/server'
 
 class Rails::ServerTest < ActiveSupport::TestCase
+  include EnvHelpers
 
   def test_environment_with_server_option
     args    = ["thin", "-e", "production"]
@@ -25,22 +27,16 @@ def test_server_option_without_environment
   end
 
   def test_environment_with_rails_env
-    rails = ENV['RAILS_ENV']
-    ENV['RAILS_ENV'] = 'production'
-    server = Rails::Server.new
-    assert_equal 'production', server.options[:environment]
-  ensure
-    ENV['RAILS_ENV'] = rails
+    with_rails_env 'production' do
+      server = Rails::Server.new
+      assert_equal 'production', server.options[:environment]
+    end
   end
 
   def test_environment_with_rack_env
-    rack, rails = ENV['RACK_ENV'], ENV['RAILS_ENV']
-    ENV['RAILS_ENV'] = nil
-    ENV['RACK_ENV'] = 'production'
-    server = Rails::Server.new
-    assert_equal 'production', server.options[:environment]
-  ensure
-    ENV['RACK_ENV'] = rack
-    ENV['RAILS_ENV'] = rails
+    with_rack_env 'production' do
+      server = Rails::Server.new
+      assert_equal 'production', server.options[:environment]
+    end
   end
 end

railties/test/env_helpers.rb

@@ -0,0 +1,26 @@
+module EnvHelpers
+  private
+
+  def with_rails_env(env)
+    switch_env 'RAILS_ENV', env do
+      switch_env 'RACK_ENV', nil do
+        yield
+      end
+    end
+  end
+
+  def with_rack_env(env)
+    switch_env 'RACK_ENV', env do
+      switch_env 'RAILS_ENV', nil do
+        yield
+      end
+    end
+  end
+
+  def switch_env(key, value)
+    old, ENV[key] = ENV[key], value
+    yield
+  ensure
+    ENV[key] = old
+  end
+end

railties/test/generators/actions_test.rb

@@ -1,8 +1,11 @@
 require 'generators/generators_test_helper'
 require 'rails/generators/rails/app/app_generator'
+require 'env_helpers'
 
 class ActionsTest < Rails::Generators::TestCase
   include GeneratorsTestHelper
+  include EnvHelpers
+
   tests Rails::Generators::AppGenerator
   arguments [destination_root]
 
@@ -154,10 +157,9 @@ def test_generate_should_run_script_generate_with_argument_and_options
 
   def test_rake_should_run_rake_command_with_default_env
     generator.expects(:run).once.with("rake log:clear RAILS_ENV=development", verbose: false)
-    old_env, ENV['RAILS_ENV'] = ENV["RAILS_ENV"], nil
-    action :rake, 'log:clear'
-  ensure
-    ENV["RAILS_ENV"] = old_env
+    with_rails_env nil do
+      action :rake, 'log:clear'
+    end
   end
 
   def test_rake_with_env_option_should_run_rake_command_in_env
@@ -167,26 +169,23 @@ def test_rake_with_env_option_should_run_rake_command_in_env
 
   def test_rake_with_rails_env_variable_should_run_rake_command_in_env
     generator.expects(:run).once.with('rake log:clear RAILS_ENV=production', verbose: false)
-    old_env, ENV["RAILS_ENV"] = ENV["RAILS_ENV"], "production"
-    action :rake, 'log:clear'
-  ensure
-    ENV["RAILS_ENV"] = old_env
+    with_rails_env "production" do
+      action :rake, 'log:clear'
+    end
   end
 
   def test_env_option_should_win_over_rails_env_variable_when_running_rake
     generator.expects(:run).once.with('rake log:clear RAILS_ENV=production', verbose: false)
-    old_env, ENV["RAILS_ENV"] = ENV["RAILS_ENV"], "staging"
-    action :rake, 'log:clear', env: 'production'
-  ensure
-    ENV["RAILS_ENV"] = old_env
+    with_rails_env "staging" do
+      action :rake, 'log:clear', env: 'production'
+    end
   end
 
   def test_rake_with_sudo_option_should_run_rake_command_with_sudo
     generator.expects(:run).once.with("sudo rake log:clear RAILS_ENV=development", verbose: false)
-    old_env, ENV['RAILS_ENV'] = ENV["RAILS_ENV"], nil
-    action :rake, 'log:clear', sudo: true
-  ensure
-    ENV["RAILS_ENV"] = old_env
+    with_rails_env nil do
+      action :rake, 'log:clear', sudo: true
+    end
   end
 
   def test_capify_should_run_the_capify_command