From 229ef2fcbde8b4d8487791f502637df23837921b Mon Sep 17 00:00:00 2001 From: Christopher Szatmary Date: Fri, 18 Sep 2020 19:34:52 -0400 Subject: [PATCH] Improve fatal and color packages (#4) * feat!: Add ability to register func to run before exiting in fatal BREAKING CHANGE: fatal.ShowStackTraces is now a function * feat: Improve color methods to strip reset values and add tests * Oops --- color/color.go | 60 ++++++++++--- color/color_test.go | 82 +++++++++++++++++ fatal/fatal.go | 79 ++++++++++++---- fatal/fatal_test.go | 213 ++++++++++++++++++++++++++++++++++++++++++++ 4 files changed, 402 insertions(+), 32 deletions(-) create mode 100644 color/color_test.go create mode 100644 fatal/fatal_test.go diff --git a/color/color.go b/color/color.go index fa48a9f..09d3e38 100644 --- a/color/color.go +++ b/color/color.go @@ -1,41 +1,75 @@ package color +import ( + "fmt" + "os" + "regexp" +) + const ( - red = "\x1b[31m" - green = "\x1b[32m" - yellow = "\x1b[33m" - blue = "\x1b[34m" - magenta = "\x1b[35m" - cyan = "\x1b[36m" - reset = "\x1b[0m" + ansiFgRed = 31 + ansiFgGreen = 32 + ansiFgYellow = 33 + ansiFgBlue = 34 + ansiFgMagenta = 35 + ansiFgCyan = 36 + asnsiFgWhite = 37 + ansiResetFg = 39 ) +// Support for NO_COLOR env var +// https://no-color.org/ +var noColor = false + +func init() { + // The standard says the value doesn't matter, only whether or not it's set + if _, ok := os.LookupEnv("NO_COLOR"); ok { + noColor = true + } +} + +func apply(str string, start, end int) string { + if noColor { + return str + } + + regex := regexp.MustCompile(fmt.Sprintf("\\x1b\\[%dm", end)) + // Remove any occurrences of reset to make sure color isn't messed up + sanitized := regex.ReplaceAllString(str, "") + return fmt.Sprintf("\x1b[%dm%s\x1b[%dm", start, sanitized, end) +} + // Red creates a red colored string func Red(str string) string { - return red + str + reset + return apply(str, ansiFgRed, ansiResetFg) } // Green creates a green colored string func Green(str string) string { - return green + str + reset + return apply(str, ansiFgGreen, ansiResetFg) } // Yellow creates a yellow colored string func Yellow(str string) string { - return yellow + str + reset + return apply(str, ansiFgYellow, ansiResetFg) } // Blue creates a blue colored string func Blue(str string) string { - return blue + str + reset + return apply(str, ansiFgBlue, ansiResetFg) } // Magenta creates a magenta colored string func Magenta(str string) string { - return magenta + str + reset + return apply(str, ansiFgMagenta, ansiResetFg) } // Cyan creates a cyan colored string func Cyan(str string) string { - return cyan + str + reset + return apply(str, ansiFgCyan, ansiResetFg) +} + +// White creates a white colored string +func White(str string) string { + return apply(str, asnsiFgWhite, ansiResetFg) } diff --git a/color/color_test.go b/color/color_test.go new file mode 100644 index 0000000..12d6adb --- /dev/null +++ b/color/color_test.go @@ -0,0 +1,82 @@ +package color + +import ( + "testing" + + "github.com/stretchr/testify/assert" +) + +func TestColors(t *testing.T) { + noColor = false + + tests := []struct { + name string + colorFn func(string) string + input string + expected string + }{ + { + "Red() test", + Red, + "foo bar", + "\x1b[31mfoo bar\x1b[39m", + }, + { + "Green() test", + Green, + "foo bar", + "\x1b[32mfoo bar\x1b[39m", + }, + { + "Yellow() test", + Yellow, + "foo bar", + "\x1b[33mfoo bar\x1b[39m", + }, + { + "Blue() test", + Blue, + "foo bar", + "\x1b[34mfoo bar\x1b[39m", + }, + { + "Magenta() test", + Magenta, + "foo bar", + "\x1b[35mfoo bar\x1b[39m", + }, + { + "Cyan() test", + Cyan, + "foo bar", + "\x1b[36mfoo bar\x1b[39m", + }, + { + "White() test", + White, + "foo bar", + "\x1b[37mfoo bar\x1b[39m", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + received := tt.colorFn(tt.input) + assert.Equal(t, tt.expected, received) + }) + } +} + +func TestStripReset(t *testing.T) { + noColor = false + + received := Red("foo \x1b[39mbar") + assert.Equal(t, "\x1b[31mfoo bar\x1b[39m", received) +} + +func TestNoColor(t *testing.T) { + noColor = true + + received := Red("foo bar") + assert.Equal(t, "foo bar", received) +} diff --git a/fatal/fatal.go b/fatal/fatal.go index 9afc09c..6703fd8 100644 --- a/fatal/fatal.go +++ b/fatal/fatal.go @@ -2,43 +2,84 @@ package fatal import ( "fmt" + "io" "os" ) -var ShowStackTraces = true +// Package state +var ( + shouldShowStackTraces = false + onExitHandler func() +) + +// Used for dependency injection in tests +// Normally having tests touch private stuff is bad +// but this is the only way I could figure out to mock os.Exit +var ( + errWriter io.Writer = os.Stderr + exitFunc func(code int) = os.Exit +) +// ShowStackTraces sets whether or not stack traces should be printed +// when ExitErr and ExitErrf are called. +func ShowStackTraces(show bool) { + shouldShowStackTraces = show +} + +// OnExit registers a handler that will run before os.Exit is called. +// This is useful for performing any clean up that would usually be called +// in a defer block since defers are not called when os.Exit is used. +func OnExit(handler func()) { + onExitHandler = handler +} + +// ExitErr prints the given message and error to stderr then exits the program. func ExitErr(err error, message string) { - fmt.Fprintln(os.Stderr, message) + fmt.Fprintln(errWriter, message) + + if err != nil { + if shouldShowStackTraces { + fmt.Fprintf(errWriter, "Error: %+v\n", err) + } else { + fmt.Fprintf(errWriter, "Error: %s\n", err) + } + } - if ShowStackTraces { - fmt.Fprintf(os.Stderr, "Error: %+v\n", err) - } else { - fmt.Fprintf(os.Stderr, "Error: %s\n", err) + if onExitHandler != nil { + onExitHandler() } - os.Exit(1) + exitFunc(1) } +// ExitErrf prints the given message and error to stderr then exits the program. +// Supports printf like formatting. func ExitErrf(err error, format string, a ...interface{}) { - fmt.Fprintf(os.Stderr, format, a...) - fmt.Fprintln(os.Stderr) + fmt.Fprintf(errWriter, format, a...) + fmt.Fprintln(errWriter) + + if err != nil { + if shouldShowStackTraces { + fmt.Fprintf(errWriter, "Error: %+v\n", err) + } else { + fmt.Fprintf(errWriter, "Error: %s\n", err) + } + } - if ShowStackTraces { - fmt.Fprintf(os.Stderr, "Error: %+v\n", err) - } else { - fmt.Fprintf(os.Stderr, "Error: %s\n", err) + if onExitHandler != nil { + onExitHandler() } - os.Exit(1) + exitFunc(1) } +// Exit prints the given message to stderr then exists the program. func Exit(message string) { - fmt.Fprintln(os.Stderr, message) - os.Exit(1) + ExitErr(nil, message) } +// Exitf prints the given message to stderr then exits the program. +// Supports printf like formatting. func Exitf(format string, a ...interface{}) { - fmt.Fprintf(os.Stderr, format, a...) - fmt.Fprintln(os.Stderr) - os.Exit(1) + ExitErrf(nil, format, a...) } diff --git a/fatal/fatal_test.go b/fatal/fatal_test.go new file mode 100644 index 0000000..d426b8f --- /dev/null +++ b/fatal/fatal_test.go @@ -0,0 +1,213 @@ +package fatal + +import ( + "bytes" + "fmt" + "regexp" + "strings" + "testing" + + "github.com/pkg/errors" + "github.com/stretchr/testify/assert" +) + +type mockExit struct { + code int +} + +func (me *mockExit) Exit(code int) { + me.code = code +} + +type client struct { + flushed bool +} + +func (c *client) Flush() { + c.flushed = true +} + +func resetState() { + // Before each: need to make sure global state is + // reset so it doesn't poison tests + ShowStackTraces(false) + OnExit(nil) +} + +func TestExit(t *testing.T) { + resetState() + + buf := &bytes.Buffer{} + me := mockExit{} + + errWriter = buf + exitFunc = me.Exit + + Exit("Something broke") + + assert.Equal(t, 1, me.code) + assert.Equal(t, "Something broke\n", buf.String()) +} + +func TestExitOnExit(t *testing.T) { + resetState() + + buf := &bytes.Buffer{} + me := mockExit{} + c := client{} + + errWriter = buf + exitFunc = me.Exit + + OnExit(func() { + c.Flush() + }) + + Exit("Something broke") + + assert.Equal(t, 1, me.code) + assert.Equal(t, "Something broke\n", buf.String()) + assert.True(t, c.flushed) +} + +func TestExitf(t *testing.T) { + resetState() + + buf := &bytes.Buffer{} + me := mockExit{} + + errWriter = buf + exitFunc = me.Exit + + Exitf("%d failures", 3) + + assert.Equal(t, 1, me.code) + assert.Equal(t, "3 failures\n", buf.String()) +} + +func TestExitfOnExit(t *testing.T) { + resetState() + + buf := &bytes.Buffer{} + me := mockExit{} + c := client{} + + errWriter = buf + exitFunc = me.Exit + + OnExit(func() { + c.Flush() + }) + + Exitf("%d failures", 3) + + assert.Equal(t, 1, me.code) + assert.Equal(t, "3 failures\n", buf.String()) + assert.True(t, c.flushed) +} + +func TestExitErr(t *testing.T) { + resetState() + + buf := &bytes.Buffer{} + me := mockExit{} + + errWriter = buf + exitFunc = me.Exit + + err := errors.New("err everything broke") + + ExitErr(err, "Something broke") + + assert.Equal(t, 1, me.code) + assert.Equal(t, "Something broke\nError: err everything broke\n", buf.String()) +} + +func TestExitErrStack(t *testing.T) { + resetState() + + buf := &bytes.Buffer{} + me := mockExit{} + + errWriter = buf + exitFunc = me.Exit + + err := errors.New("err everything broke") + + ShowStackTraces(true) + + ExitErr(err, "Something broke") + + assert.Equal(t, 1, me.code) + + expected := "Something broke\n" + + "Error: err everything broke\n" + + "github.com/TouchBistro/goutils/fatal.TestExitErrStack\n" + + "\t.+" + testFormatRegexp(t, 0, err, buf.String(), expected) +} + +func TestExitErrf(t *testing.T) { + resetState() + + buf := &bytes.Buffer{} + me := mockExit{} + + errWriter = buf + exitFunc = me.Exit + + err := errors.New("err everything broke") + + ExitErrf(err, "%d failures", 3) + + assert.Equal(t, 1, me.code) + assert.Equal(t, "3 failures\nError: err everything broke\n", buf.String()) +} + +func TestExitErrStackf(t *testing.T) { + resetState() + + buf := &bytes.Buffer{} + me := mockExit{} + + errWriter = buf + exitFunc = me.Exit + + err := errors.New("err everything broke") + + ShowStackTraces(true) + + ExitErrf(err, "%d failures", 3) + + assert.Equal(t, 1, me.code) + + expected := "3 failures\n" + + "Error: err everything broke\n" + + "github.com/TouchBistro/goutils/fatal.TestExitErrStack\n" + + "\t.+" + testFormatRegexp(t, 0, err, buf.String(), expected) +} + +// Taken from https://github.com/pkg/errors/blob/614d223910a179a466c1767a985424175c39b465/format_test.go#L387 +// Helper to test string with regexp +func testFormatRegexp(t *testing.T, n int, arg interface{}, format, want string) { + t.Helper() + got := fmt.Sprintf(format, arg) + gotLines := strings.SplitN(got, "\n", -1) + wantLines := strings.SplitN(want, "\n", -1) + + if len(wantLines) > len(gotLines) { + t.Errorf("test %d: wantLines(%d) > gotLines(%d):\n got: %q\nwant: %q", n+1, len(wantLines), len(gotLines), got, want) + return + } + + for i, w := range wantLines { + match, err := regexp.MatchString(w, gotLines[i]) + if err != nil { + t.Fatal(err) + } + if !match { + t.Errorf("test %d: line %d: fmt.Sprintf(%q, err):\n got: %q\nwant: %q", n+1, i+1, format, got, want) + } + } +}