diff options
Diffstat (limited to 'cli/internal/logstreamer')
| -rw-r--r-- | cli/internal/logstreamer/logstreamer.go | 159 | ||||
| -rw-r--r-- | cli/internal/logstreamer/logstreamer_test.go | 114 |
2 files changed, 273 insertions, 0 deletions
diff --git a/cli/internal/logstreamer/logstreamer.go b/cli/internal/logstreamer/logstreamer.go new file mode 100644 index 0000000..4379c25 --- /dev/null +++ b/cli/internal/logstreamer/logstreamer.go @@ -0,0 +1,159 @@ +// Copyright (c) 2013 Kevin van Zonneveld <kevin@vanzonneveld.net>. All rights reserved. +// Source: https://github.com/kvz/logstreamer +// SPDX-License-Identifier: MIT +package logstreamer + +import ( + "bytes" + "io" + "log" + "os" + "strings" +) + +type Logstreamer struct { + Logger *log.Logger + buf *bytes.Buffer + // If prefix == stdout, colors green + // If prefix == stderr, colors red + // Else, prefix is taken as-is, and prepended to anything + // you throw at Write() + prefix string + // if true, saves output in memory + record bool + persist string + + // Adds color to stdout & stderr if terminal supports it + colorOkay string + colorFail string + colorReset string +} + +func NewLogstreamer(logger *log.Logger, prefix string, record bool) *Logstreamer { + streamer := &Logstreamer{ + Logger: logger, + buf: bytes.NewBuffer([]byte("")), + prefix: prefix, + record: record, + persist: "", + colorOkay: "", + colorFail: "", + colorReset: "", + } + + if strings.HasPrefix(os.Getenv("TERM"), "xterm") { + streamer.colorOkay = "\x1b[32m" + streamer.colorFail = "\x1b[31m" + streamer.colorReset = "\x1b[0m" + } + + return streamer +} + +func (l *Logstreamer) Write(p []byte) (n int, err error) { + if n, err = l.buf.Write(p); err != nil { + return + } + + err = l.OutputLines() + return +} + +func (l *Logstreamer) Close() error { + if err := l.Flush(); err != nil { + return err + } + l.buf = bytes.NewBuffer([]byte("")) + return nil +} + +func (l *Logstreamer) Flush() error { + p := make([]byte, l.buf.Len()) + if _, err := l.buf.Read(p); err != nil { + return err + } + + l.out(string(p)) + return nil +} + +func (l *Logstreamer) OutputLines() error { + for { + line, err := l.buf.ReadString('\n') + + if len(line) > 0 { + if strings.HasSuffix(line, "\n") { + l.out(line) + } else { + // put back into buffer, it's not a complete line yet + // Close() or Flush() have to be used to flush out + // the last remaining line if it does not end with a newline + if _, err := l.buf.WriteString(line); err != nil { + return err + } + } + } + + if err == io.EOF { + break + } + + if err != nil { + return err + } + } + + return nil +} + +func (l *Logstreamer) FlushRecord() string { + buffer := l.persist + l.persist = "" + return buffer +} + +func (l *Logstreamer) out(str string) { + if len(str) < 1 { + return + } + + if l.record { + l.persist = l.persist + str + } + + if l.prefix == "stdout" { + str = l.colorOkay + l.prefix + l.colorReset + " " + str + } else if l.prefix == "stderr" { + str = l.colorFail + l.prefix + l.colorReset + " " + str + } + + l.Logger.Print(str) +} + +// PrettyStdoutWriter wraps an ioWriter so it can add string +// prefixes to every message it writes to stdout. +type PrettyStdoutWriter struct { + w io.Writer + Prefix string +} + +var _ io.Writer = (*PrettyStdoutWriter)(nil) + +// NewPrettyStdoutWriter returns an instance of PrettyStdoutWriter +func NewPrettyStdoutWriter(prefix string) *PrettyStdoutWriter { + return &PrettyStdoutWriter{ + w: os.Stdout, + Prefix: prefix, + } +} + +func (psw *PrettyStdoutWriter) Write(p []byte) (int, error) { + str := psw.Prefix + string(p) + n, err := psw.w.Write([]byte(str)) + + if err != nil { + return n, err + } + + return len(p), nil +} diff --git a/cli/internal/logstreamer/logstreamer_test.go b/cli/internal/logstreamer/logstreamer_test.go new file mode 100644 index 0000000..94d8a82 --- /dev/null +++ b/cli/internal/logstreamer/logstreamer_test.go @@ -0,0 +1,114 @@ +// Copyright (c) 2013 Kevin van Zonneveld <kevin@vanzonneveld.net>. All rights reserved. +// Source: https://github.com/kvz/logstreamer +// SPDX-License-Identifier: MIT +package logstreamer + +import ( + "bufio" + "bytes" + "fmt" + "log" + "os" + "os/exec" + "strings" + "testing" +) + +func TestLogstreamerOk(t *testing.T) { + // Create a logger (your app probably already has one) + logger := log.New(os.Stdout, "--> ", log.Ldate|log.Ltime) + + // Setup a streamer that we'll pipe cmd.Stdout to + logStreamerOut := NewLogstreamer(logger, "stdout", false) + defer logStreamerOut.Close() + // Setup a streamer that we'll pipe cmd.Stderr to. + // We want to record/buffer anything that's written to this (3rd argument true) + logStreamerErr := NewLogstreamer(logger, "stderr", true) + defer logStreamerErr.Close() + + // Execute something that succeeds + cmd := exec.Command( + "ls", + "-al", + ) + cmd.Stderr = logStreamerErr + cmd.Stdout = logStreamerOut + + // Reset any error we recorded + logStreamerErr.FlushRecord() + + // Execute command + err := cmd.Start() + + // Failed to spawn? + if err != nil { + t.Fatal("ERROR could not spawn command.", err.Error()) + } + + // Failed to execute? + err = cmd.Wait() + if err != nil { + t.Fatal("ERROR command finished with error. ", err.Error(), logStreamerErr.FlushRecord()) + } +} + +func TestLogstreamerErr(t *testing.T) { + // Create a logger (your app probably already has one) + logger := log.New(os.Stdout, "--> ", log.Ldate|log.Ltime) + + // Setup a streamer that we'll pipe cmd.Stdout to + logStreamerOut := NewLogstreamer(logger, "stdout", false) + defer logStreamerOut.Close() + // Setup a streamer that we'll pipe cmd.Stderr to. + // We want to record/buffer anything that's written to this (3rd argument true) + logStreamerErr := NewLogstreamer(logger, "stderr", true) + defer logStreamerErr.Close() + + // Execute something that succeeds + cmd := exec.Command( + "ls", + "nonexisting", + ) + cmd.Stderr = logStreamerErr + cmd.Stdout = logStreamerOut + + // Reset any error we recorded + logStreamerErr.FlushRecord() + + // Execute command + err := cmd.Start() + + // Failed to spawn? + if err != nil { + logger.Print("ERROR could not spawn command. ") + } + + // Failed to execute? + err = cmd.Wait() + if err != nil { + fmt.Printf("Good. command finished with %s. %s. \n", err.Error(), logStreamerErr.FlushRecord()) + } else { + t.Fatal("This command should have failed") + } +} + +func TestLogstreamerFlush(t *testing.T) { + const text = "Text without newline" + + var buffer bytes.Buffer + byteWriter := bufio.NewWriter(&buffer) + + logger := log.New(byteWriter, "", 0) + logStreamerOut := NewLogstreamer(logger, "", false) + defer logStreamerOut.Close() + + logStreamerOut.Write([]byte(text)) + logStreamerOut.Flush() + byteWriter.Flush() + + s := strings.TrimSpace(buffer.String()) + + if s != text { + t.Fatalf("Expected '%s', got '%s'.", text, s) + } +} |
