Add support for Rails console

lib/rails/app_env/console.rb

@@ -0,0 +1,25 @@
+require "rails/commands/console/irb_console"
+
+module Rails
+  module AppEnv
+    class Console < Rails::Console::IRBConsole
+      def colorized_env
+        return super if Rails.env == Rails.app_env
+        super + ":" + colorized_app_env
+      end
+
+      private
+
+      def colorized_app_env
+        case Rails.app_env
+        when "development"
+          IRB::Color.colorize("dev", [:MAGENTA])
+        when "production"
+          IRB::Color.colorize("prod", [:RED])
+        else
+          IRB::Color.colorize(Rails.app_env, [:MAGENTA])
+        end
+      end
+    end
+  end
+end

lib/rails/app_env/railtie.rb

@@ -15,6 +15,13 @@ class Railtie < Rails::Railtie
           Rails.app_env
         end
       end
+
+      console do |app|
+        require_relative "console"
+
+        app.config.console = Rails::AppEnv::Console.new(app)
+        puts "Loading #{Rails.app_env} application environment (rails-app_env #{Rails::AppEnv::VERSION})" # standard:disable Rails/Output
+      end
     end
   end
 end

test/features/console_helpers.rb

@@ -0,0 +1,62 @@
+require "pty"
+
+module ConsoleHelpers
+  private
+
+  DUMMY_RAILS = File.expand_path "../../dummy/bin/rails", __FILE__
+
+  def with_pty
+    PTY.open do |primary, replica|
+      @primary, @replica = primary, replica
+
+      yield
+    end
+  end
+
+  def spawn_console(options, wait_for_prompt: true, env: {})
+    # Test should not depend on user's irbrc file
+    home_tmp_dir = Dir.mktmpdir
+
+    env = env.with_defaults(
+      "TERM" => "dumb",
+      "HOME" => home_tmp_dir,
+      "APP_ENV" => nil,
+      "RAILS_ENV" => nil,
+      "SECRET_KEY_BASE_DUMMY" => "1"
+    )
+
+    pid = Process.spawn(
+      env,
+      "#{DUMMY_RAILS} console #{options}",
+      in: @replica, out: @replica, err: @replica
+    )
+
+    if wait_for_prompt
+      assert_output "> ", @primary
+    end
+
+    pid
+  ensure
+    FileUtils.remove_entry(home_tmp_dir)
+  end
+
+  def write_prompt(command, expected_output = nil, prompt: "> ")
+    @primary.puts command.to_s
+    assert_output command, @primary
+    assert_output expected_output, @primary if expected_output
+    assert_output prompt, @primary
+  end
+
+  def assert_output(expected, io, timeout = 5)
+    timeout = Time.current + timeout
+
+    output = +""
+    until output.include?(expected) || Time.current > timeout
+      if IO.select([io], [], [], 0.1)
+        output << io.read(1)
+      end
+    end
+
+    assert_includes output, expected, "#{expected.inspect} expected, but got:\n\n#{output}"
+  end
+end

test/features/console_test.rb

@@ -0,0 +1,244 @@
+require_relative "../test_helper"
+require_relative "console_helpers"
+
+module Rails::AppEnv::FeaturesTest
+  class ConsoleTest < ActiveSupport::TestCase
+    include ConsoleHelpers
+
+    TEST_CASES = {
+      "includes Rails.app_env when APP_ENV is custom and RAILS_ENV is custom": {
+        expected: "bar:foo",
+        env: {
+          "APP_ENV" => "foo",
+          "RAILS_ENV" => "bar"
+        }
+      },
+      "includes Rails.app_env when APP_ENV is custom and RAILS_ENV is production": {
+        expected: "prod:foo",
+        env: {
+          "APP_ENV" => "foo",
+          "RAILS_ENV" => "production"
+        }
+      },
+      "includes Rails.app_env when APP_ENV is custom and RAILS_ENV is development": {
+        expected: "dev:foo",
+        env: {
+          "APP_ENV" => "foo",
+          "RAILS_ENV" => "development"
+        }
+      },
+      "includes Rails.app_env when APP_ENV is custom and RAILS_ENV is test": {
+        expected: "test:foo",
+        env: {
+          "APP_ENV" => "foo",
+          "RAILS_ENV" => "test"
+        }
+      },
+
+      "includes Rails.app_env when APP_ENV is production and RAILS_ENV is custom": {
+        expected: "bar:prod",
+        env: {
+          "APP_ENV" => "production",
+          "RAILS_ENV" => "bar"
+        }
+      },
+      "includes Rails.app_env when APP_ENV is production and RAILS_ENV is development": {
+        expected: "dev:prod",
+        env: {
+          "APP_ENV" => "production",
+          "RAILS_ENV" => "development"
+        }
+      },
+      "includes Rails.app_env when APP_ENV is production and RAILS_ENV is test": {
+        expected: "test:prod",
+        env: {
+          "APP_ENV" => "production",
+          "RAILS_ENV" => "test"
+        }
+      },
+
+      "includes Rails.app_env when APP_ENV is development and RAILS_ENV is custom": {
+        expected: "bar:dev",
+        env: {
+          "APP_ENV" => "development",
+          "RAILS_ENV" => "bar"
+        }
+      },
+      "includes Rails.app_env when APP_ENV is development and RAILS_ENV is production": {
+        expected: "prod:dev",
+        env: {
+          "APP_ENV" => "development",
+          "RAILS_ENV" => "production"
+        }
+      },
+      "includes Rails.app_env when APP_ENV is development and RAILS_ENV is test": {
+        expected: "test:dev",
+        env: {
+          "APP_ENV" => "development",
+          "RAILS_ENV" => "test"
+        }
+      },
+
+      "does not include Rails.app_env when both APP_ENV and RAILS_ENV are custom": {
+        expected: "foo",
+        env: {
+          "APP_ENV" => "foo",
+          "RAILS_ENV" => "foo"
+        }
+      },
+      "does not include Rails.app_env when both APP_ENV and RAILS_ENV are production": {
+        expected: "prod",
+        env: {
+          "APP_ENV" => "production",
+          "RAILS_ENV" => "production"
+        }
+      },
+      "does not include Rails.app_env when both APP_ENV and RAILS_ENV are development": {
+        expected: "dev",
+        env: {
+          "APP_ENV" => "development",
+          "RAILS_ENV" => "development"
+        }
+      },
+      "does not include Rails.app_env when both APP_ENV and RAILS_ENV are test": {
+        expected: "test",
+        env: {
+          "APP_ENV" => "test",
+          "RAILS_ENV" => "test"
+        }
+      },
+
+      "does not include Rails.app_env when APP_ENV is blank and RAILS_ENV is custom": {
+        expected: "foo",
+        env: {
+          "RAILS_ENV" => "foo"
+        }
+      },
+      "does not include Rails.app_env when APP_ENV is blank and RAILS_ENV is production": {
+        expected: "prod",
+        env: {
+          "RAILS_ENV" => "production"
+        }
+      },
+      "does not include Rails.app_env when APP_ENV is blank and RAILS_ENV is development": {
+        expected: "dev",
+        env: {
+          "RAILS_ENV" => "development"
+        }
+      },
+      "does not include Rails.app_env when APP_ENV is blank and RAILS_ENV is test": {
+        expected: "test",
+        env: {
+          "RAILS_ENV" => "test"
+        }
+      },
+
+      "includes Rails.app_env when APP_ENV is custom and RAILS_ENV is blank": {
+        expected: "dev:foo",
+        env: {
+          "APP_ENV" => "foo"
+        }
+      },
+      "includes Rails.app_env when APP_ENV is production and RAILS_ENV is blank": {
+        expected: "dev:prod",
+        env: {
+          "APP_ENV" => "production"
+        }
+      },
+      "includes Rails.app_env when APP_ENV is test and RAILS_ENV is blank": {
+        expected: "dev:test",
+        env: {
+          "APP_ENV" => "test"
+        }
+      },
+      "does not include Rails.app_env when APP_ENV is development and RAILS_ENV is blank": {
+        expected: "dev",
+        env: {
+          "APP_ENV" => "development"
+        }
+      },
+
+      "does not include Rails.app_env when both APP_ENV and RAILS_ENV are blank": {
+        expected: "dev"
+      }
+    }
+
+    BANNER_TEST_CASES = {
+      "when APP_ENV is custom and RAILS_ENV is custom": {
+        expected: "foo",
+        env: {
+          "APP_ENV" => "foo",
+          "RAILS_ENV" => "bar"
+        }
+      },
+      "when both APP_ENV and RAILS_ENV are production": {
+        expected: "production",
+        env: {
+          "APP_ENV" => "production",
+          "RAILS_ENV" => "production"
+        }
+      },
+      "when both APP_ENV and RAILS_ENV are development": {
+        expected: "development",
+        env: {
+          "APP_ENV" => "development",
+          "RAILS_ENV" => "development"
+        }
+      },
+      "when APP_ENV is present and RAILS_ENV is production": {
+        expected: "foo",
+        env: {
+          "APP_ENV" => "foo",
+          "RAILS_ENV" => "production"
+        }
+      },
+      "when APP_ENV is present and RAILS_ENV is blank": {
+        expected: "foo",
+        env: {
+          "APP_ENV" => "foo"
+        }
+      },
+      "when APP_ENV is blank and RAILS_ENV is present": {
+        expected: "foo",
+        env: {
+          "RAILS_ENV" => "foo"
+        }
+      },
+      "when both APP_ENV and RAILS_ENV are blank": {
+        expected: "development" # Default Rails environment
+      }
+    }
+
+    TEST_CASES.each_key do |name|
+      class_eval <<~RUBY, __FILE__, __LINE__ + 1
+        test "Rails console prompt #{name}" do
+          assert_rails_console_prompt(**TEST_CASES[:"#{name}"])
+        end
+      RUBY
+    end
+
+    BANNER_TEST_CASES.each_key do |name|
+      class_eval <<~RUBY, __FILE__, __LINE__ + 1
+        test "Rails console prints Rails.app_env and gem version on start #{name}" do
+          assert_rails_console_banner(**BANNER_TEST_CASES[:"#{name}"])
+        end
+      RUBY
+    end
+
+    private
+
+    def assert_rails_console_prompt(expected:, env: {})
+      with_pty do
+        spawn_console("--sandbox", env: env, wait_for_prompt: true)
+        write_prompt "123", prompt: "dummy(#{expected})> "
+      end
+    end
+
+    def assert_rails_console_banner(expected:, env: {})
+      with_pty do
+        spawn_console("--sandbox", env: env, wait_for_prompt: false)
+        assert_output "Loading #{expected} application environment (rails-app_env #{Rails::AppEnv::VERSION})", @primary
+      end
+    end
+  end
+end

test/units/app_env/console_test.rb

@@ -0,0 +1,51 @@
+require "minitest/mock"
+require_relative "../../test_helper"
+require "rails/app_env/console"
+
+class Rails::AppEnv::ConsoleTest < ActiveSupport::TestCase
+  test "Console is a kind of Rails::Console::IRBConsole" do
+    assert_kind_of Rails::Console::IRBConsole, Rails::AppEnv::Console.new(nil)
+  end
+
+  test "Console#colorized_env prints Rails.env only when Rails.app_env is same as Rails.env" do
+    Rails.stub :app_env, "foo" do
+      Rails.stub :env, "foo" do
+        console = Rails::AppEnv::Console.new(nil)
+        assert_equal_without_color "foo", console.colorized_env
+      end
+    end
+  end
+
+  test "Console#colorized_env prints {Rails.env:Rails.app_env} when Rails.app_env is diff from Rails.env" do
+    Rails.stub :app_env, "foo" do
+      Rails.stub :env, "bar" do
+        console = Rails::AppEnv::Console.new(nil)
+        assert_equal_without_color "bar:foo", console.colorized_env
+      end
+    end
+  end
+
+  test "Console#colorized_env shorten Rails.app_env to dev when Rails.app_env is production" do
+    Rails.stub :app_env, "production" do
+      Rails.stub :env, "bar" do
+        console = Rails::AppEnv::Console.new(nil)
+        assert_equal_without_color "bar:prod", console.colorized_env
+      end
+    end
+  end
+
+  test "Console#colorized_env shorten Rails.app_env to dev when Rails.app_env is development" do
+    Rails.stub :app_env, "development" do
+      Rails.stub :env, "bar" do
+        console = Rails::AppEnv::Console.new(nil)
+        assert_equal_without_color "bar:dev", console.colorized_env
+      end
+    end
+  end
+
+  private
+
+  def assert_equal_without_color(expected, actual)
+    assert_equal expected, actual.gsub(/\e\[(\d+)(;\d+)*m/, "")
+  end
+end