Change default output & add WithTTY option

output.go

@@ -26,12 +26,13 @@ type Output struct {
 	tty     io.Writer
 	environ Environ
 
-	unsafe  bool
-	cache   bool
-	fgSync  *sync.Once
-	fgColor Color
-	bgSync  *sync.Once
-	bgColor Color
+	assumeTTY bool
+	unsafe    bool
+	cache     bool
+	fgSync    *sync.Once
+	fgColor   Color
+	bgSync    *sync.Once
+	bgColor   Color
 }
 
 // Environ is an interface for getting environment variables.
@@ -55,6 +56,11 @@ func DefaultOutput() *Output {
 	return output
 }
 
+// SetDefaultOutput sets the default global output.
+func SetDefaultOutput(o *Output) {
+	output = o
+}
+
 // NewOutput returns a new Output for the given file descriptor.
 func NewOutput(tty io.Writer, opts ...OutputOption) *Output {
 	o := &Output{
@@ -80,21 +86,21 @@ func NewOutput(tty io.Writer, opts ...OutputOption) *Output {
 	return o
 }
 
-// WithEnvironment returns a new Output for the given environment.
+// WithEnvironment returns a new OutputOption for the given environment.
 func WithEnvironment(environ Environ) OutputOption {
 	return func(o *Output) {
 		o.environ = environ
 	}
 }
 
-// WithProfile returns a new Output for the given profile.
+// WithProfile returns a new OutputOption for the given profile.
 func WithProfile(profile Profile) OutputOption {
 	return func(o *Output) {
 		o.Profile = profile
 	}
 }
 
-// WithColorCache returns a new Output with fore- and background color values
+// WithColorCache returns a new OutputOption with fore- and background color values
 // pre-fetched and cached.
 func WithColorCache(v bool) OutputOption {
 	return func(o *Output) {
@@ -106,9 +112,19 @@ func WithColorCache(v bool) OutputOption {
 	}
 }
 
-// WithUnsafe returns a new Output with unsafe mode enabled. Unsafe mode doesn't
+// WithTTY returns a new OutputOption to assume whether or not the output is a TTY.
+// This is useful when mocking console output.
+func WithTTY(v bool) OutputOption {
+	return func(o *Output) {
+		o.assumeTTY = v
+	}
+}
+
+// WithUnsafe returns a new OutputOption with unsafe mode enabled. Unsafe mode doesn't
 // check whether or not the terminal is a TTY.
 //
+// This option supersedes WithTTY.
+//
 // This is useful when mocking console output and enforcing ANSI escape output
 // e.g. on SSH sessions.
 func WithUnsafe() OutputOption {

termenv.go

@@ -25,7 +25,7 @@ const (
 )
 
 func (o *Output) isTTY() bool {
-	if o.unsafe {
+	if o.assumeTTY || o.unsafe {
 		return true
 	}
 	if len(o.environ.Getenv("CI")) > 0 {

termenv_test.go

@@ -4,6 +4,7 @@ import (
 	"bytes"
 	"fmt"
 	"image/color"
+	"io/ioutil"
 	"os"
 	"strings"
 	"testing"
@@ -405,3 +406,12 @@ func TestEnableVirtualTerminalProcessing(t *testing.T) {
 		t.Fatalf("expected <nil>, got %v", err)
 	}
 }
+
+func TestWithTTY(t *testing.T) {
+	for _, v := range []bool{true, false} {
+		o := NewOutput(ioutil.Discard, WithTTY(v))
+		if o.isTTY() != v {
+			t.Fatalf("expected WithTTY(%t) to set isTTY to %t", v, v)
+		}
+	}
+}