github.com/onflow/flow-go@v0.35.7-crescendo-preview.23-atree-inlining/utils/unittest/logging.go (about)

     1  package unittest
     2  
     3  import (
     4  	"flag"
     5  	"fmt"
     6  	"io"
     7  	"os"
     8  	"path/filepath"
     9  	"runtime"
    10  	"strings"
    11  	"sync"
    12  	"testing"
    13  	"time"
    14  
    15  	"github.com/rs/zerolog"
    16  )
    17  
    18  var verbose = flag.Bool("vv", false, "print debugging logs")
    19  var globalOnce sync.Once
    20  
    21  func LogVerbose() {
    22  	*verbose = true
    23  }
    24  
    25  // Logger returns a zerolog
    26  // use -vv flag to print debugging logs for tests
    27  func Logger() zerolog.Logger {
    28  	writer := io.Discard
    29  	if *verbose {
    30  		writer = os.Stderr
    31  	}
    32  
    33  	return LoggerWithWriterAndLevel(writer, zerolog.InfoLevel)
    34  }
    35  
    36  func LoggerWithWriterAndLevel(writer io.Writer, level zerolog.Level) zerolog.Logger {
    37  	globalOnce.Do(func() {
    38  		zerolog.TimestampFunc = func() time.Time { return time.Now().UTC() }
    39  	})
    40  	log := zerolog.New(writer).Level(level).With().Timestamp().Logger()
    41  	return log
    42  }
    43  
    44  // go:noinline
    45  func LoggerForTest(t *testing.T, level zerolog.Level) zerolog.Logger {
    46  	_, file, _, ok := runtime.Caller(1)
    47  	if !ok {
    48  		file = "???"
    49  	}
    50  	return LoggerWithLevel(level).With().
    51  		Str("testfile", filepath.Base(file)).Str("testcase", t.Name()).Logger()
    52  }
    53  
    54  func LoggerWithLevel(level zerolog.Level) zerolog.Logger {
    55  	return LoggerWithWriterAndLevel(os.Stderr, level)
    56  }
    57  
    58  func LoggerWithName(mod string) zerolog.Logger {
    59  	return Logger().With().Str("module", mod).Logger()
    60  }
    61  
    62  func NewLoggerHook() LoggerHook {
    63  	return LoggerHook{
    64  		logs: &strings.Builder{},
    65  		mu:   &sync.Mutex{},
    66  	}
    67  }
    68  
    69  func HookedLogger() (zerolog.Logger, LoggerHook) {
    70  	hook := NewLoggerHook()
    71  	log := zerolog.New(io.Discard).Hook(hook)
    72  	return log, hook
    73  }
    74  
    75  // LoggerHook implements the zerolog.Hook interface and can be used to capture
    76  // logs for testing purposes.
    77  type LoggerHook struct {
    78  	logs *strings.Builder
    79  	mu   *sync.Mutex
    80  }
    81  
    82  // Logs returns the logs as a string
    83  func (hook LoggerHook) Logs() string {
    84  	hook.mu.Lock()
    85  	defer hook.mu.Unlock()
    86  
    87  	return hook.logs.String()
    88  }
    89  
    90  // Run implements zerolog.Hook and appends the log message to the log.
    91  func (hook LoggerHook) Run(_ *zerolog.Event, level zerolog.Level, msg string) {
    92  	// for tests that need to test logger.Fatal(), this is useful because the parent test process will read from stdout
    93  	// to determine if the test sub-process (that generated the logger.Fatal() call) called logger.Fatal() with the expected message
    94  	if level == zerolog.FatalLevel {
    95  		fmt.Println(msg)
    96  	}
    97  
    98  	hook.mu.Lock()
    99  	defer hook.mu.Unlock()
   100  
   101  	hook.logs.WriteString(msg)
   102  }