Fix various bugs with console arguments.

Allow hyphens in environment names again.

5 files changed, 180 insertions, 125 deletions

diff --git a/railties/lib/rails/commands.rb b/railties/lib/rails/commands.rb
index 7f473c237c..8816387d34 100644
--- a/railties/lib/rails/commands.rb
+++ b/railties/lib/rails/commands.rb
@@ -36,9 +36,17 @@ when 'benchmarker', 'profiler'
 
 when 'console'
   require 'rails/commands/console'
+  options = Rails::Console.parse_arguments(ARGV)
+
+  # RAILS_ENV needs to be set before config/application is required
+  ENV['RAILS_ENV'] = options[:environment] if options[:environment]
+
+  # shift ARGV so IRB doesn't freak
+  ARGV.shift if ARGV.first && ARGV.first[0] != '-'
+
   require APP_PATH
   Rails.application.require_environment!
-  Rails::Console.start(Rails.application)
+  Rails::Console.start(Rails.application, options)
 
 when 'server'
   # Change to the application's path if there is no config.ru file in current dir.
diff --git a/railties/lib/rails/commands/console.rb b/railties/lib/rails/commands/console.rb
index b95df3e545..92cee6b638 100644
--- a/railties/lib/rails/commands/console.rb
+++ b/railties/lib/rails/commands/console.rb
@@ -4,21 +4,12 @@ require 'irb/completion'
 
 module Rails
   class Console
-    attr_reader :options, :app, :console, :arguments
-
-    def self.start(*args)
-      new(*args).start
-    end
-
-    def initialize(app, arguments = ARGV)
-      @app       = app
-      @arguments = arguments
-      app.load_console
-      @console   = app.config.console || IRB
-    end
+    class << self
+      def start(*args)
+        new(*args).start
+      end
 
-    def options
-      @options ||= begin
+      def parse_arguments(arguments)
         options = {}
 
         OptionParser.new do |opt|
@@ -31,20 +22,38 @@ module Rails
           opt.parse!(arguments)
         end
 
+        if arguments.first && arguments.first[0] != '-'
+          env = arguments.first
+          options[:environment] = %w(production development test).detect {|e| e =~ /^#{env}/} || env
+        end
+
         options
       end
     end
 
+    attr_reader :options, :app, :console
+
+    def initialize(app, options={})
+      @app     = app
+      @options = options
+      app.load_console
+      @console = app.config.console || IRB
+    end
+
     def sandbox?
       options[:sandbox]
     end
 
+    def environment
+      options[:environment] ||= ENV['RAILS_ENV'] || 'development'
+    end
+
     def environment?
-      options[:environment]
+      environment
     end
 
     def set_environment!
-      Rails.env = options[:environment]
+      Rails.env = environment
     end
 
     def debugger?
@@ -53,9 +62,7 @@ module Rails
 
     def start
       app.sandbox = sandbox?
-
       require_debugger if debugger?
-
       set_environment! if environment?
 
       if sandbox?
@@ -82,8 +89,3 @@ module Rails
     end
   end
 end
-
-# Has to set the RAILS_ENV before config/application is required
-if ARGV.first && !ARGV.first.index("-") && env = ARGV.shift # has to shift the env ARGV so IRB doesn't freak
-  ENV['RAILS_ENV'] = %w(production development test).detect {|e| e =~ /^#{env}/} || env
-end
diff --git a/railties/lib/rails/commands/dbconsole.rb b/railties/lib/rails/commands/dbconsole.rb
index cc7caffc3d..cc0552184a 100644
--- a/railties/lib/rails/commands/dbconsole.rb
+++ b/railties/lib/rails/commands/dbconsole.rb
@@ -5,66 +5,19 @@ require 'rbconfig'
 
 module Rails
   class DBConsole
-    attr_reader :arguments, :config
-
+    attr_reader :config, :arguments
+    
     def self.start
-      new(config).start
+      new.start
     end
 
-    def self.config
-      config = begin
-        YAML.load(ERB.new(IO.read("config/database.yml")).result)
-      rescue SyntaxError, StandardError
-        require APP_PATH
-        Rails.application.config.database_configuration
-      end
-
-      unless config[env]
-        abort "No database is configured for the environment '#{env}'"
-      end
-
-      config[env]
-    end
-
-    def self.env
-      if Rails.respond_to?(:env)
-        Rails.env
-      else
-        ENV["RAILS_ENV"] || ENV["RACK_ENV"] || "development"
-      end
-    end
-
-    def initialize(config, arguments = ARGV)
-      @config, @arguments = config, arguments
+    def initialize(arguments = ARGV)
+      @arguments = arguments
     end
 
     def start
-      include_password = false
-      options = {}
-      OptionParser.new do |opt|
-        opt.banner = "Usage: rails dbconsole [environment] [options]"
-        opt.on("-p", "--include-password", "Automatically provide the password from database.yml") do |v|
-          include_password = true
-        end
-
-        opt.on("--mode [MODE]", ['html', 'list', 'line', 'column'],
-          "Automatically put the sqlite3 database in the specified mode (html, list, line, column).") do |mode|
-            options['mode'] = mode
-        end
-
-        opt.on("--header") do |h|
-          options['header'] = h
-        end
-
-        opt.on("-h", "--help", "Show this help message.") do
-          puts opt
-          exit
-        end
-
-        opt.parse!(arguments)
-        abort opt.to_s unless (0..1).include?(arguments.size)
-      end
-
+      options = parse_arguments(arguments)
+      ENV['RAILS_ENV'] = options[:environment] || environment
 
       case config["adapter"]
       when /^mysql/
@@ -76,7 +29,7 @@ module Rails
           'encoding'  => '--default-character-set'
         }.map { |opt, arg| "#{arg}=#{config[opt]}" if config[opt] }.compact
 
-        if config['password'] && include_password
+        if config['password'] && options['include_password']
           args << "--password=#{config['password']}"
         elsif config['password'] && !config['password'].to_s.empty?
           args << "-p"
@@ -90,7 +43,7 @@ module Rails
         ENV['PGUSER']     = config["username"] if config["username"]
         ENV['PGHOST']     = config["host"] if config["host"]
         ENV['PGPORT']     = config["port"].to_s if config["port"]
-        ENV['PGPASSWORD'] = config["password"].to_s if config["password"] && include_password
+        ENV['PGPASSWORD'] = config["password"].to_s if config["password"] && options['include_password']
         find_cmd_and_exec('psql', config["database"])
 
       when "sqlite"
@@ -110,7 +63,7 @@ module Rails
 
         if config['username']
           logon = config['username']
-          logon << "/#{config['password']}" if config['password'] && include_password
+          logon << "/#{config['password']}" if config['password'] && options['include_password']
           logon << "@#{config['database']}" if config['database']
         end
 
@@ -121,8 +74,73 @@ module Rails
       end
     end
 
+    def config
+      @config ||= begin
+        cfg = begin
+          cfg = YAML.load(ERB.new(IO.read("config/database.yml")).result)
+        rescue SyntaxError, StandardError
+          require APP_PATH
+          Rails.application.config.database_configuration
+        end
+
+        unless cfg[environment]
+          abort "No database is configured for the environment '#{environment}'"
+        end
+
+        cfg[environment]
+      end
+    end
+
+    def environment
+      if Rails.respond_to?(:env)
+        Rails.env
+      else
+        ENV["RAILS_ENV"] || ENV["RACK_ENV"] || "development"
+      end
+    end
+
     protected
 
+    def parse_arguments(arguments)
+      options = {}
+      
+      OptionParser.new do |opt|
+        opt.banner = "Usage: rails dbconsole [environment] [options]"
+        opt.on("-p", "--include-password", "Automatically provide the password from database.yml") do |v|
+          options['include_password'] = true
+        end
+
+        opt.on("--mode [MODE]", ['html', 'list', 'line', 'column'],
+          "Automatically put the sqlite3 database in the specified mode (html, list, line, column).") do |mode|
+            options['mode'] = mode
+        end
+
+        opt.on("--header") do |h|
+          options['header'] = h
+        end
+        
+        opt.on("-h", "--help", "Show this help message.") do
+          puts opt
+          exit
+        end
+
+        opt.on("-e", "--environment=name", String,
+          "Specifies the environment to run this console under (test/development/production).",
+          "Default: development"
+        ) { |v| options[:environment] = v.strip }
+
+        opt.parse!(arguments)
+        abort opt.to_s unless (0..1).include?(arguments.size)
+      end
+
+      if arguments.first && arguments.first[0] != '-'
+        env = arguments.first
+        options[:environment] = %w(production development test).detect {|e| e =~ /^#{env}/} || env
+      end
+      
+      options
+    end
+
     def find_cmd_and_exec(commands, *args)
       commands = Array(commands)
 
@@ -145,8 +163,3 @@ module Rails
     end
   end
 end
-
-# Has to set the RAILS_ENV before config/application is required
-if ARGV.first && !ARGV.first.index("-") && env = ARGV.first
-  ENV['RAILS_ENV'] = %w(production development test).detect {|e| e =~ /^#{env}/} || env
-end
diff --git a/railties/test/commands/console_test.rb b/railties/test/commands/console_test.rb
index 78648a16b3..91ede1cb68 100644
--- a/railties/test/commands/console_test.rb
+++ b/railties/test/commands/console_test.rb
@@ -3,47 +3,44 @@ require 'rails/commands/console'
 
 class Rails::ConsoleTest < ActiveSupport::TestCase
   class FakeConsole
+    def self.start; end
   end
 
   def setup
   end
 
   def test_sandbox_option
-    console = Rails::Console.new(app, ["--sandbox"])
+    console = Rails::Console.new(app, parse_arguments(["--sandbox"]))
     assert console.sandbox?
   end
 
   def test_short_version_of_sandbox_option
-    console = Rails::Console.new(app, ["-s"])
+    console = Rails::Console.new(app, parse_arguments(["-s"]))
     assert console.sandbox?
   end
 
   def test_debugger_option
-    console = Rails::Console.new(app, ["--debugger"])
+    console = Rails::Console.new(app, parse_arguments(["--debugger"]))
     assert console.debugger?
   end
 
   def test_no_options
-    console = Rails::Console.new(app, [])
+    console = Rails::Console.new(app, parse_arguments([]))
     assert !console.debugger?
     assert !console.sandbox?
   end
 
   def test_start
-    app.expects(:sandbox=).with(nil)
     FakeConsole.expects(:start)
-
     start
-
     assert_match(/Loading \w+ environment \(Rails/, output)
   end
 
   def test_start_with_debugger
-    app.expects(:sandbox=).with(nil)
+    rails_console = Rails::Console.new(app, parse_arguments(["--debugger"]))
     rails_console.expects(:require_debugger).returns(nil)
-    FakeConsole.expects(:start)
 
-    start ["--debugger"]
+    silence_stream(STDOUT) { rails_console.start }
   end
 
   def test_start_with_sandbox
@@ -56,42 +53,63 @@ class Rails::ConsoleTest < ActiveSupport::TestCase
   end
 
   def test_console_with_environment
-    app.expects(:sandbox=).with(nil)
-    FakeConsole.expects(:start)
-
     start ["-e production"]
+    assert_match(/\sproduction\s/, output)
+  end
+
+  def test_console_defaults_to_IRB
+    config = mock("config", :console => nil)
+    app = mock("app", :config => config)
+    app.expects(:load_console).returns(nil)
 
-    assert_match(/production/, output)
+    assert_equal IRB, Rails::Console.new(app).console
   end
 
-  def test_console_with_rails_environment
-    app.expects(:sandbox=).with(nil)
-    FakeConsole.expects(:start)
+  def test_default_environment_with_no_rails_env
+    with_rails_env nil do
+      start
+      assert_match /\sdevelopment\s/, output
+    end
+  end
 
-    start ["RAILS_ENV=production"]
+  def test_default_environment_with_rails_env
+    with_rails_env 'special-production' do
+      start
+      assert_match /\sspecial-production\s/, output
+    end
+  end
+  
+  def test_e_option
+    start ['-e', 'special-production']
+    assert_match /\sspecial-production\s/, output
+  end
 
-    assert_match(/production/, output)
+  def test_environment_option
+    start ['--environment=special-production']
+    assert_match /\sspecial-production\s/, output
   end
 
+  def test_rails_env_is_production_when_first_argument_is_p
+    start ['p']
+    assert_match /\sproduction\s/, output
+  end
 
-  def test_console_defaults_to_IRB
-    config = mock("config", :console => nil)
-    app = mock("app", :config => config)
-    app.expects(:load_console).returns(nil)
+  def test_rails_env_is_test_when_first_argument_is_t
+    start ['t']
+    assert_match /\stest\s/, output
+  end
 
-    assert_equal IRB, Rails::Console.new(app).console
+  def test_rails_env_is_development_when_argument_is_d
+    start ['d']
+    assert_match /\sdevelopment\s/, output
   end
 
   private
 
   attr_reader :output
 
-  def rails_console
-    @rails_console ||= Rails::Console.new(app)
-  end
-
   def start(argv = [])
-    rails_console.stubs(:arguments => argv)
+    rails_console = Rails::Console.new(app, parse_arguments(argv))
     @output = output = capture(:stdout) { rails_console.start }
   end
 
@@ -99,8 +117,21 @@ class Rails::ConsoleTest < ActiveSupport::TestCase
     @app ||= begin
       config = mock("config", :console => FakeConsole)
       app = mock("app", :config => config)
+      app.stubs(:sandbox=).returns(nil)
       app.expects(:load_console)
       app
     end
   end
+
+  def parse_arguments(args)
+    Rails::Console.parse_arguments(args)
+  end
+
+  def with_rails_env(env)
+    original_rails_env = ENV['RAILS_ENV']
+    ENV['RAILS_ENV'] = env
+    yield
+  ensure
+    ENV['RAILS_ENV'] = original_rails_env
+  end
 end
diff --git a/railties/test/commands/dbconsole_test.rb b/railties/test/commands/dbconsole_test.rb
index 562b83713b..1b477eb98a 100644
--- a/railties/test/commands/dbconsole_test.rb
+++ b/railties/test/commands/dbconsole_test.rb
@@ -10,35 +10,36 @@ class Rails::DBConsoleTest < ActiveSupport::TestCase
     Rails::DBConsole.const_set(:APP_PATH, "erb")
 
     app_config({})
-    capture_abort { Rails::DBConsole.config }
+    capture_abort { Rails::DBConsole.new.config }
     assert aborted
     assert_match(/No database is configured for the environment '\w+'/, output)
 
     app_config(test: "with_init")
-    assert_equal Rails::DBConsole.config, "with_init"
+    assert_equal Rails::DBConsole.new.config, "with_init"
 
     app_db_file("test:\n  without_init")
-    assert_equal Rails::DBConsole.config, "without_init"
+    assert_equal Rails::DBConsole.new.config, "without_init"
 
     app_db_file("test:\n  <%= Rails.something_app_specific %>")
-    assert_equal Rails::DBConsole.config, "with_init"
+    assert_equal Rails::DBConsole.new.config, "with_init"
 
     app_db_file("test:\n\ninvalid")
-    assert_equal Rails::DBConsole.config, "with_init"
+    assert_equal Rails::DBConsole.new.config, "with_init"
   end
 
   def test_env
-    assert_equal Rails::DBConsole.env, "test"
+    assert_equal Rails::DBConsole.new.environment, "test"
+
+    ENV['RAILS_ENV'] = nil
 
     Rails.stubs(:respond_to?).with(:env).returns(false)
-    assert_equal Rails::DBConsole.env, "test"
+    assert_equal Rails::DBConsole.new.environment, "development"
 
-    ENV['RAILS_ENV'] = nil
     ENV['RACK_ENV'] = "rack_env"
-    assert_equal Rails::DBConsole.env, "rack_env"
+    assert_equal Rails::DBConsole.new.environment, "rack_env"
 
     ENV['RAILS_ENV'] = "rails_env"
-    assert_equal Rails::DBConsole.env, "rails_env"
+    assert_equal Rails::DBConsole.new.environment, "rails_env"
   ensure
     ENV['RAILS_ENV'] = "test"
   end