github.com/anth0d/nomad@v0.0.0-20221214183521-ae3a0a2cad06/helper/testlog/testlog.go (about)

     1  // Package testlog creates a *log.Logger backed by *testing.T to ease logging
     2  // in tests. This allows logs from components being tested to only be printed
     3  // if the test fails (or the verbose flag is specified).
     4  package testlog
     5  
     6  import (
     7  	"bytes"
     8  	"fmt"
     9  	"io"
    10  	"log"
    11  	"os"
    12  
    13  	hclog "github.com/hashicorp/go-hclog"
    14  )
    15  
    16  // LogPrinter is the methods of testing.T (or testing.B) needed by the test
    17  // logger.
    18  type LogPrinter interface {
    19  	Logf(format string, args ...interface{})
    20  }
    21  
    22  // NewWriter creates a new io.Writer backed by a Logger.
    23  func NewWriter(t LogPrinter) io.Writer {
    24  	return os.Stderr
    25  }
    26  
    27  // NewPrefixWriter creates a new io.Writer backed by a Logger with a custom
    28  // prefix per Write.
    29  func NewPrefixWriter(t LogPrinter, prefix string) io.Writer {
    30  	return &prefixStderr{[]byte(prefix)}
    31  }
    32  
    33  // New returns a new test logger. See https://golang.org/pkg/log/#New
    34  func New(t LogPrinter, prefix string, flag int) *log.Logger {
    35  	return log.New(os.Stderr, prefix, flag)
    36  }
    37  
    38  // WithPrefix returns a new test logger with the Lmicroseconds flag set.
    39  func WithPrefix(t LogPrinter, prefix string) *log.Logger {
    40  	return New(t, prefix, log.Lmicroseconds)
    41  }
    42  
    43  // Logger returns a new test logger with the Lmicroseconds flag set and no prefix.
    44  //
    45  // Note: only use this where HCLogger cannot be used (i.e. RPC yamux configuration).
    46  func Logger(t LogPrinter) *log.Logger {
    47  	return WithPrefix(t, "")
    48  }
    49  
    50  // HCLogger returns a new test hc-logger.
    51  //
    52  // Default log level is TRACE. Set NOMAD_TEST_LOG_LEVEL for custom log level.
    53  func HCLogger(t LogPrinter) hclog.InterceptLogger {
    54  	logger, _ := HCLoggerNode(t, -1)
    55  	return logger
    56  }
    57  
    58  // HCLoggerTestLevel returns the level in which hc log should emit logs.
    59  //
    60  // Default log level is TRACE. Set NOMAD_TEST_LOG_LEVEL for custom log level.
    61  func HCLoggerTestLevel() hclog.Level {
    62  	level := hclog.Trace
    63  	envLogLevel := os.Getenv("NOMAD_TEST_LOG_LEVEL")
    64  	if envLogLevel != "" {
    65  		level = hclog.LevelFromString(envLogLevel)
    66  	}
    67  	return level
    68  }
    69  
    70  // HCLoggerNode returns a new hc-logger, but with a prefix indicating the node number
    71  // on each log line. Useful for TestServer in tests with more than one server.
    72  //
    73  // Default log level is TRACE. Set NOMAD_TEST_LOG_LEVEL for custom log level.
    74  func HCLoggerNode(t LogPrinter, node int32) (hclog.InterceptLogger, io.Writer) {
    75  	var output io.Writer = os.Stderr
    76  	if node > -1 {
    77  		output = NewPrefixWriter(t, fmt.Sprintf("node-%03d ", node))
    78  	}
    79  	opts := &hclog.LoggerOptions{
    80  		Level:           HCLoggerTestLevel(),
    81  		Output:          output,
    82  		IncludeLocation: true,
    83  	}
    84  	return hclog.NewInterceptLogger(opts), output
    85  }
    86  
    87  type prefixStderr struct {
    88  	prefix []byte
    89  }
    90  
    91  // Write to stdout with a prefix per call containing non-whitespace characters.
    92  func (w *prefixStderr) Write(p []byte) (int, error) {
    93  	if len(p) == 0 {
    94  		return 0, nil
    95  	}
    96  
    97  	// Skip prefix if only writing whitespace
    98  	if len(bytes.TrimSpace(p)) == 0 {
    99  		return os.Stderr.Write(p)
   100  	}
   101  
   102  	// decrease likely hood of partial line writes that may mess up test
   103  	// indicator success detection
   104  	buf := make([]byte, 0, len(w.prefix)+len(p))
   105  	buf = append(buf, w.prefix...)
   106  	buf = append(buf, p...)
   107  
   108  	return os.Stderr.Write(buf)
   109  }