Add helper for testing TTY output.

3 files changed, 111 insertions, 10 deletions

diff --git a/Library/Homebrew/test/cask/cli/search_spec.rb b/Library/Homebrew/test/cask/cli/search_spec.rb
index c382614a4..4c779cae3 100644
--- a/Library/Homebrew/test/cask/cli/search_spec.rb
+++ b/Library/Homebrew/test/cask/cli/search_spec.rb
@@ -1,12 +1,12 @@
 describe Hbc::CLI::Search, :cask do
   before(:each) do
-    allow($stdout).to receive(:tty?).and_return(true)
+    allow(Tty).to receive(:width).and_return(0)
   end
 
   it "lists the available Casks that match the search term" do
     expect {
       Hbc::CLI::Search.run("local")
-    }.to output(<<-EOS.undent).to_stdout
+    }.to output(<<-EOS.undent).to_stdout.as_tty
       ==> Partial Matches
       local-caffeine
       local-transmission
@@ -16,49 +16,49 @@ describe Hbc::CLI::Search, :cask do
   it "shows that there are no Casks matching a search term that did not result in anything" do
     expect {
       Hbc::CLI::Search.run("foo-bar-baz")
-    }.to output("No Cask found for \"foo-bar-baz\".\n").to_stdout
+    }.to output("No Cask found for \"foo-bar-baz\".\n").to_stdout.as_tty
   end
 
   it "lists all available Casks with no search term" do
     expect {
       Hbc::CLI::Search.run
-    }.to output(/local-caffeine/).to_stdout
+    }.to output(/local-caffeine/).to_stdout.as_tty
   end
 
   it "ignores hyphens in search terms" do
     expect {
       Hbc::CLI::Search.run("lo-cal-caffeine")
-    }.to output(/local-caffeine/).to_stdout
+    }.to output(/local-caffeine/).to_stdout.as_tty
   end
 
   it "ignores hyphens in Cask tokens" do
     expect {
       Hbc::CLI::Search.run("localcaffeine")
-    }.to output(/local-caffeine/).to_stdout
+    }.to output(/local-caffeine/).to_stdout.as_tty
   end
 
   it "accepts multiple arguments" do
     expect {
       Hbc::CLI::Search.run("local caffeine")
-    }.to output(/local-caffeine/).to_stdout
+    }.to output(/local-caffeine/).to_stdout.as_tty
   end
 
   it "accepts a regexp argument" do
     expect {
       Hbc::CLI::Search.run("/^local-c[a-z]ffeine$/")
-    }.to output("==> Regexp Matches\nlocal-caffeine\n").to_stdout
+    }.to output("==> Regexp Matches\nlocal-caffeine\n").to_stdout.as_tty
   end
 
   it "Returns both exact and partial matches" do
     expect {
       Hbc::CLI::Search.run("test-opera")
-    }.to output(/^==> Exact Match\ntest-opera\n==> Partial Matches\ntest-opera-mail/).to_stdout
+    }.to output(/^==> Exact Match\ntest-opera\n==> Partial Matches\ntest-opera-mail/).to_stdout.as_tty
   end
 
   it "does not search the Tap name" do
     expect {
       Hbc::CLI::Search.run("caskroom")
-    }.to output(/^No Cask found for "caskroom"\.\n/).to_stdout
+    }.to output(/^No Cask found for "caskroom"\.\n/).to_stdout.as_tty
   end
 
   it "doesn't highlight packages that aren't installed" do
diff --git a/Library/Homebrew/test/spec_helper.rb b/Library/Homebrew/test/spec_helper.rb
index 75540caad..d04ca0088 100644
--- a/Library/Homebrew/test/spec_helper.rb
+++ b/Library/Homebrew/test/spec_helper.rb
@@ -23,6 +23,7 @@ require "test/support/helper/shutup"
 require "test/support/helper/fixtures"
 require "test/support/helper/formula"
 require "test/support/helper/mktmpdir"
+require "test/support/helper/output_as_tty"
 require "test/support/helper/rubocop"
 
 require "test/support/helper/spec/shared_context/homebrew_cask" if OS.mac?
@@ -45,6 +46,7 @@ RSpec.configure do |config|
   config.include(Test::Helper::Fixtures)
   config.include(Test::Helper::Formula)
   config.include(Test::Helper::MkTmpDir)
+  config.include(Test::Helper::OutputAsTTY)
   config.include(Test::Helper::RuboCop)
 
   config.before(:each, :needs_compat) do
diff --git a/Library/Homebrew/test/support/helper/output_as_tty.rb b/Library/Homebrew/test/support/helper/output_as_tty.rb
new file mode 100644
index 000000000..bd72a8f3f
--- /dev/null
+++ b/Library/Homebrew/test/support/helper/output_as_tty.rb
@@ -0,0 +1,99 @@
+require "delegate"
+
+module Test
+  module Helper
+    module OutputAsTTY
+      module TTYTrue
+        def tty?
+          true
+        end
+
+        alias isatty tty?
+      end
+
+      # This is a custom wrapper for the `output` matcher,
+      # used for testing output to a TTY:
+      #
+      #   expect {
+      #     print "test" if $stdout.tty?
+      #   }.to output("test").to_stdout.as_tty
+      #
+      #   expect {
+      #     # command
+      #   }.to output(...).to_stderr.as_tty.with_color
+      #
+      class Output < SimpleDelegator
+        def matches?(block)
+          return super(block) unless @tty
+
+          colored_tty_block = if @output == :stdout
+            lambda do
+              $stdout.extend(TTYTrue)
+              block.call
+            end
+          elsif @output == :stderr
+            lambda do
+              $stderr.extend(TTYTrue)
+              block.call
+            end
+          else
+            raise "`as_tty` can only be chained to `stdout` or `stderr`."
+          end
+
+          return super(colored_tty_block) if @colors
+
+          uncolored_tty_block = lambda do
+            begin
+              out_stream = StringIO.new
+              err_stream = StringIO.new
+
+              old_stdout = $stdout
+              old_stderr = $stderr
+
+              $stdout = out_stream
+              $stderr = err_stream
+
+              colored_tty_block.call
+            ensure
+              $stdout = old_stdout
+              $stderr = old_stderr
+
+              $stdout.print Tty.strip_ansi(out_stream.string)
+              $stderr.print Tty.strip_ansi(err_stream.string)
+            end
+          end
+
+          super(uncolored_tty_block)
+        end
+
+        def to_stdout
+          @output = :stdout
+          super
+          self
+        end
+
+        def to_stderr
+          @output = :stderr
+          super
+          self
+        end
+
+        def as_tty
+          @tty = true
+          self
+        end
+
+        def with_color
+          @colors = true
+          return self if @tty
+          raise "`with_color` can only be chained to `as_tty`."
+        end
+      end
+
+      def output(*args)
+        core_matcher = super(*args)
+        Output.new(core_matcher)
+      end
+    end
+  end
+end