Restore color-strip behavior of Console class

lib/airbrussh/console.rb

@@ -4,6 +4,10 @@
 module Airbrussh
   # Helper class that wraps an IO object and provides methods for truncating
   # output, assuming the IO object represents a console window.
+  #
+  # If color is disabled for the IO object (based on Airbrussh::Configuration),
+  # any ANSI color codes will also be stripped from the output.
+  #
   class Console
     attr_reader :output, :config
 
@@ -13,15 +17,19 @@ def initialize(output, config=Airbrussh.configuration)
     end
 
     # Writes to the IO after first truncating the output to fit the console
-    # width. A newline is always added.
+    # width. If color is disabled for the underlying IO, ANSI colors are also
+    # removed from the output. A newline is always added.
     def print_line(obj="")
       string = obj.to_s
+
       string = truncate_to_console_width(string) if console_width
+      string = strip_ascii_color(string) unless color_enabled?
+
       write(string + "\n")
       output.flush
     end
 
-    # Writes directly through to the IO with no truncation logic.
+    # Writes directly through to the IO with no truncation or color logic.
     # No newline is added.
     def write(string)
       output.write(string || "")
@@ -36,8 +44,7 @@ def truncate_to_console_width(string)
       if strip_ascii_color(string).length > width
         width -= ellipsis.length
         string.chop! while strip_ascii_color(string).length > width
-        string << ellipsis
-        terminate_color(string)
+        string << "#{ellipsis}\e[0m"
       else
         string
       end
@@ -58,8 +65,8 @@ def console_width
 
     private
 
-    def terminate_color(string)
-      string << ("\e[0m" if contains_ascii_color?(string)).to_s
+    def color_enabled?
+      @color_enabled ||= config.colors(output).enabled?
     end
 
     def contains_ascii_color?(string)

test/airbrussh/console_test.rb

@@ -10,6 +10,47 @@ def setup
     @output = StringIO.new
   end
 
+  def test_color_is_allowed_for_tty
+    console = configured_console(:tty => true) do |config|
+      config.color = :auto
+    end
+    console.print_line("The \e[0;32;49mgreen\e[0m text")
+    assert_equal("The \e[0;32;49mgreen\e[0m text\n", output)
+  end
+
+  def test_color_is_can_be_forced
+    console = configured_console(:tty => false) do |config|
+      config.color = true
+    end
+    console.print_line("The \e[0;32;49mgreen\e[0m text")
+    assert_equal("The \e[0;32;49mgreen\e[0m text\n", output)
+  end
+
+  def test_color_is_can_be_forced_via_env
+    console = configured_console(:tty => false) do |config|
+      config.color = :auto
+    end
+    ENV.stubs(:[]).with("SSHKIT_COLOR").returns("1")
+    console.print_line("The \e[0;32;49mgreen\e[0m text")
+    assert_equal("The \e[0;32;49mgreen\e[0m text\n", output)
+  end
+
+  def test_color_is_stripped_for_non_tty
+    console = configured_console(:tty => false) do |config|
+      config.color = :auto
+    end
+    console.print_line("The \e[0;32;49mgreen\e[0m text")
+    assert_equal("The green text\n", output)
+  end
+
+  def test_color_can_be_disabled_for_tty
+    console = configured_console(:tty => true) do |config|
+      config.color = false
+    end
+    console.print_line("The \e[0;32;49mgreen\e[0m text")
+    assert_equal("The green text\n", output)
+  end
+
   def test_truncates_to_winsize
     console = configured_console(:tty => true) do |config|
       config.color = false