github.com/hernad/nomad@v1.6.112/helper/testlog/testlog.go (about)

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